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

[KVOffload][P2P] Generic P2P secondary tier: peer lookup and serving via ParentManager (#48021)

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 Jul 26, 2026 at 11:45 UTC da3a252fd13f51c22657bfc8650936f2fbb5b6f3

11 files changed +2673 -396

docs/features/kv_offloading_usage.md

+66 -1

@@ -157,7 +157,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti

157

158 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.

159

160 -PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes.

160 +The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected.

161

163 | --- | --- | --- | --- |

@@ -176,6 +176,71 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on

176 - `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide.

177 - `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision.

178

179 +#### Orchestration-Layer Protocol

180 +

181 +The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations.

182 +

183 +Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes.

184 +

185 +Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache.

186 +

187 +Each key names the **remote counterpart** this peer transfers with (not this

188 +peer's own role), so the name reads as "the remote ___ I transfer with".

189 +

191 +| --- | --- | --- | --- |

195 +

196 +Field semantics:

197 +

198 +- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID.

199 +- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)).

200 +- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank.

201 +

202 +Allowed and forbidden combinations:

203 +

204 +- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder.

205 +- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together.

206 +

207 +Minimal examples (values that would appear in the request's `kv_transfer_params`):

208 +

209 +```python

210 +# Prefill producer — compute and keep KV for a remote decoder to pull

211 +kv_transfer_params = {"remote_decoder": {"kv_request_id": "<unique-transfer-id>"}}

212 +

213 +# Decode consumer — pull KV from a specific prefiller (classic P/D)

214 +kv_transfer_params = {

215 + "remote_prefiller": {

216 + "kv_request_id": "<unique-transfer-id>",

217 + "remote_host": "<prefiller-node-ip>",

218 + "remote_port": 5710,

219 + }

220 +}

221 +

222 +# P2P consumer — pull whatever the source already has cached

223 +kv_transfer_params = {

224 + "remote_kv_source": {

225 + "kv_request_id": "<unique-transfer-id>",

226 + "remote_host": "<source-node-ip>",

227 + "remote_port": 5710,

228 + }

229 +}

230 +```

231 +

232 +Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above:

233 +

234 +1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)).

235 +2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step.

236 +3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes.

237 +4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits.

238 +5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes).

239 +6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status.

240 +7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine.

241 +

242 +In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached.

243 +

244 ## Tuning Tips

245

246 - `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/p2p_connector_proxy.py

+22 -7

@@ -81,6 +81,8 @@ async def lifespan(app: FastAPI):

81 app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size))

83 mode = "decoder-first" if global_args.decoder_first else "prefiller-first"

84 + if global_args.p2p:

85 + mode += ",p2p"

86 pd_host = global_args.p2p_connector_host

87 pd_port = global_args.p2p_connector_port

88 n_pref = len(app.state.prefill_clients)

@@ -145,6 +147,14 @@ def parse_args():

147 help="Send decode request before prefill so decoder is already "

148 "waiting when KV blocks arrive (decoder-first mode)",

149 )

150 + p.add_argument(

151 + "--p2p",

152 + action="store_true",

153 + help="P2P mode: do not inject kv_transfer_params on the prefiller; "

154 + "on the decoder, inject kv_transfer_params with a top-level "

155 + "'remote_kv_source' block ({kv_request_id, remote_host, remote_port}) "

156 + "instead of the default 'remote_prefiller' block.",

157 + )

158 p.add_argument(

159 "--prefiller-dp-size",

160 type=int,

@@ -206,11 +216,14 @@ def _auth_headers(request_id: str) -> dict:

216 async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None):

217 """Send a prefill-only request (max_tokens=1) to the prefiller."""

218 data = req_data.copy()

209 - data["kv_transfer_params"] = {

210 - "decode": {

211 - "kv_request_id": request_id,

212 - },

213 - }

219 + if global_args.p2p:

220 + data.pop("kv_transfer_params", None)

221 + else:

222 + data["kv_transfer_params"] = {

223 + "remote_decoder": {

224 + "kv_request_id": request_id,

225 + },

226 + }

227 data["stream"] = False

228 data["max_tokens"] = 1

229 data.pop("max_completion_tokens", None)

@@ -259,8 +272,9 @@ async def _handle_completions(api: str, request: Request):

272 # Inject the prefiller's P2PConnector address so the decoder can pull

273 # KV blocks from it. remote_port = base + prefill_rank targets the

274 # replica that produced the KV (base+0 == base when dp=1).

275 + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller"

276 req_data["kv_transfer_params"] = {

263 - "prefill": {

277 + decoder_key: {

278 "kv_request_id": request_id,

279 "remote_host": global_args.p2p_connector_host,

280 "remote_port": global_args.p2p_connector_port + prefill_rank,

@@ -311,9 +325,10 @@ async def _handle_completions_decoder_first(api: str, request: Request):

325 prefill_client = _get_next(request.app, "prefill")

326 decode_client = _get_next(request.app, "decode")

327

328 + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller"

329 decode_data = req_data.copy()

330 decode_data["kv_transfer_params"] = {

316 - "prefill": {

331 + decoder_key: {

332 "kv_request_id": request_id,

333 "remote_host": global_args.p2p_connector_host,

334 "remote_port": global_args.p2p_connector_port + prefill_rank,

tests/v1/kv_offload/tiering/p2p/test_manager.py

+256 -51

@@ -21,9 +21,11 @@ from vllm.v1.kv_offload.tiering.p2p import manager as manager_module

21 from vllm.v1.kv_offload.tiering.p2p.manager import (

22 _UNBOUND_STORE_TIMEOUT_S,

23 P2PSecondaryTierManager,

24 + _annotate_req_context,

25 )

26 from vllm.v1.kv_offload.tiering.p2p.session import (

27 LoadResult,

28 + SessionCloseResult,

29 SessionPollResult,

30 StoreResult,

31 )

@@ -33,15 +35,15 @@ from vllm.v1.kv_offload.tiering.p2p.session import (

35 # ---------------------------------------------------------------------------

36 -def _prefill_kv_params(

38 +def _remote_prefiller_kv_params(

39 remote_host: str = "10.0.0.1",

40 remote_port: int = 8000,

41 kv_request_id: str = "req-1",

42 ) -> dict:

41 - """Decoder-side kv_transfer_params: ``prefill`` sub-dict carries

43 + """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries

44 kv_request_id + remote_host + remote_port."""

45 return {

44 - "prefill": {

46 + "remote_prefiller": {

47 "kv_request_id": kv_request_id,

48 "remote_host": remote_host,

49 "remote_port": remote_port,

@@ -49,14 +51,34 @@ def _prefill_kv_params(

51 }

52 -def _decode_kv_params(kv_request_id: str = "req-1") -> dict:

53 - """Prefiller-side kv_transfer_params: ``decode`` sub-dict carries

54 +def _remote_kv_source_kv_params(

55 + remote_host: str = "10.0.0.1",

56 + remote_port: int = 8000,

57 + kv_request_id: str = "req-1",

58 +) -> dict:

59 + """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has

60 + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port)."""

61 + return {

62 + "remote_kv_source": {

63 + "kv_request_id": kv_request_id,

64 + "remote_host": remote_host,

65 + "remote_port": remote_port,

66 + },

67 + }

68 +

69 +

70 +def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict:

71 + """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries

72 kv_request_id only."""

55 - return {"decode": {"kv_request_id": kv_request_id}}

73 + return {"remote_decoder": {"kv_request_id": kv_request_id}}

76 def _req_context(kv_params: dict | None = None) -> ReqContext:

59 - return ReqContext(req_id="test", kv_transfer_params=kv_params)

77 + ctx = ReqContext(req_id="test", kv_transfer_params=kv_params)

78 + # Mirror on_new_request: parse the P2P routing state once and cache it,

79 + # so lookup/submit_*/on_request_finished can read it back via get_state.

80 + _annotate_req_context(ctx)

81 + return ctx

84 def _job_metadata(

@@ -82,38 +104,78 @@ def _make_manager() -> P2PSecondaryTierManager:

104 """Create a manager with stubbed __init__."""

105 mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager)

106 mgr._local_id = "127.0.0.1:7777"

107 + mgr._hash_seed = "0"

108 mgr._finished_jobs = []

109 mgr._failed_req_ids = set()

110 mgr._sessions = {}

111 mgr._kv_to_session = {}

112 mgr._unbound_stores = {}

113 + mgr._failed_serve_ctxs = []

114 return mgr

115

116

117 +def _init_offloading_spec() -> SimpleNamespace:

118 + """Minimal offloading_spec for driving the real __init__."""

119 + return SimpleNamespace(

120 + config=SimpleNamespace(parallel=SimpleNamespace(data_parallel_index=0)),

121 + blocks_per_chunk=1,

122 + )

123 +

124 +

125 +# ---------------------------------------------------------------------------

126 +# Tests for __init__ PYTHONHASHSEED assertion

127 +# ---------------------------------------------------------------------------

128 +

129 +

130 +class TestInitHashSeedAssertion:

131 + def test_missing_pythonhashseed_raises(self, monkeypatch):

132 + """P2P instance refuses to start when PYTHONHASHSEED is unset."""

133 + monkeypatch.delenv("PYTHONHASHSEED", raising=False)

134 + with pytest.raises(ValueError, match="PYTHONHASHSEED"):

135 + P2PSecondaryTierManager(

136 + offloading_spec=_init_offloading_spec(),

137 + primary_kv_view=memoryview(bytearray(16)),

138 + )

139 +

140 + def test_pythonhashseed_set_succeeds(self, monkeypatch):

141 + """With PYTHONHASHSEED set, __init__ records it for the handshake."""

142 + monkeypatch.setenv("PYTHONHASHSEED", "12345")

143 + monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object())

144 + monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object())

145 + monkeypatch.setattr(

146 + manager_module.FileMapper,

147 + "from_offloading_spec",

148 + lambda **k: SimpleNamespace(get_run_config=lambda: {}),

149 + )

150 + mgr = P2PSecondaryTierManager(

151 + offloading_spec=_init_offloading_spec(),

152 + primary_kv_view=memoryview(bytearray(16)),

153 + )

154 + assert mgr._hash_seed == "12345"

155 +

156 +

157 # ---------------------------------------------------------------------------

94 -# Tests for _remote_id_from_params

158 +# Tests for _peer_id_from_params

159 # ---------------------------------------------------------------------------

160

161

98 -class TestRemoteIdFromParams:

162 +class TestPeerIdFromParams:

163 def test_valid_params(self):

100 - result = P2PSecondaryTierManager._remote_id_from_params(

164 + result = manager_module._peer_id_from_params(

165 {"remote_host": "10.0.0.1", "remote_port": 8000}

166 )

167 assert result == "10.0.0.1:8000"

168

169 def test_missing_host(self):

106 - result = P2PSecondaryTierManager._remote_id_from_params({"remote_port": 8000})

170 + result = manager_module._peer_id_from_params({"remote_port": 8000})

171 assert result is None

172

173 def test_missing_port(self):

110 - result = P2PSecondaryTierManager._remote_id_from_params(

111 - {"remote_host": "10.0.0.1"}

112 - )

174 + result = manager_module._peer_id_from_params({"remote_host": "10.0.0.1"})

175 assert result is None

176

177 def test_empty_dict(self):

116 - result = P2PSecondaryTierManager._remote_id_from_params({})

178 + result = manager_module._peer_id_from_params({})

179 assert result is None

180

181

@@ -130,35 +192,105 @@ class TestLookup:

192

193 def test_lookup_returns_miss_without_required_fields(self):

194 mgr = _make_manager()

133 - ctx = _req_context(kv_params={"prefill": {"remote_host": "x"}})

195 + ctx = _req_context(kv_params={"remote_prefiller": {"remote_host": "x"}})

196 assert mgr.lookup(b"key", ctx) is LookupResult.MISS

197

198 def test_lookup_returns_hit_for_valid_request(self):

199 mgr = _make_manager()

138 - ctx = _req_context(kv_params=_prefill_kv_params())

200 + ctx = _req_context(kv_params=_remote_prefiller_kv_params())

201 assert mgr.lookup(b"key", ctx) is LookupResult.HIT

202

203 def test_lookup_returns_miss_for_failed_request(self):

204 mgr = _make_manager()

205 mgr._failed_req_ids.add("req-1")

144 - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))

206 + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1"))

207 assert mgr.lookup(b"key", ctx) is LookupResult.MISS

208

209 def test_lookup_returns_hit_for_different_request_id(self):

210 mgr = _make_manager()

211 mgr._failed_req_ids.add("req-1")

150 - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-2"))

212 + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-2"))

213 assert mgr.lookup(b"key", ctx) is LookupResult.HIT

214

215 def test_lookup_returns_miss_without_prefill_key(self):

154 - """No ``prefill`` sub-dict means the request was not routed for

216 + """No ``remote_prefiller`` sub-dict means the request was not routed for

217 remote prefill — local prefill should run instead, so lookup()

156 - returns MISS even when a stale ``decode`` block is present."""

218 + returns MISS even when a stale ``remote_decoder`` block is present."""

219 mgr = _make_manager()

158 - ctx = _req_context(kv_params=_decode_kv_params())

220 + ctx = _req_context(kv_params=_remote_decoder_kv_params())

221 assert mgr.lookup(b"key", ctx) is LookupResult.MISS

222

223

224 +# ---------------------------------------------------------------------------

225 +# Tests for serve_external_requests

226 +# ---------------------------------------------------------------------------

227 +

228 +

229 +class _RecordingParent:

230 + """Minimal ParentManager stub recording on_request_finished calls."""

231 +

232 + def __init__(self) -> None:

233 + self.finished: list[str] = []

234 +

235 + def on_new_request(self, ctx):

236 + from vllm.v1.kv_offload.base import RequestOffloadingContext

237 +

238 + return RequestOffloadingContext()

239 +

240 + def lookup(self, key, ctx):

241 + return LookupResult.MISS

242 +

243 + def create_store_job(self, keys, ctx):

244 + raise AssertionError("unreachable")

245 +

246 + def on_request_finished(self, ctx) -> None:

247 + self.finished.append(ctx.req_id)

248 +

249 +

250 +class _RecordingSession:

251 + """Fake P2PSession that records the parent it was served with."""

252 +

253 + def __init__(self) -> None:

254 + self.served_with: list[object] = []

255 +

256 + def serve_external_requests(self, parent) -> None:

257 + self.served_with.append(parent)

258 +

259 +

260 +class TestServeExternalRequests:

261 + def test_flushes_failed_serve_ctxs_then_serves_each_session(self):

262 + """serve_external_requests releases the failed serves left by

263 + reaped sessions via parent.on_request_finished (clearing the

264 + queue), then delegates to every live session with the same parent."""

265 + mgr = _make_manager()

266 + ctx = ReqContext(req_id="p2p:peer:req-1:lu1")

267 + mgr._failed_serve_ctxs = [ctx]

268 + sess_a = _RecordingSession()

269 + sess_b = _RecordingSession()

270 + mgr._sessions = {"a": sess_a, "b": sess_b} # type: ignore[assignment]

271 +

272 + parent = _RecordingParent()

273 + mgr.serve_external_requests(parent) # type: ignore[arg-type]

274 +

275 + # Failed serve released and queue cleared.

276 + assert parent.finished == ["p2p:peer:req-1:lu1"]

277 + assert mgr._failed_serve_ctxs == []

278 + # Every live session served with the same parent handle.

279 + assert sess_a.served_with == [parent]

280 + assert sess_b.served_with == [parent]

281 +

282 + def test_no_failed_serves_still_serves_sessions(self):

283 + mgr = _make_manager()

284 + sess = _RecordingSession()

285 + mgr._sessions = {"a": sess} # type: ignore[assignment]

286 +

287 + parent = _RecordingParent()

288 + mgr.serve_external_requests(parent) # type: ignore[arg-type]

289 +

290 + assert parent.finished == []

291 + assert sess.served_with == [parent]

292 +

293 +

294 # ---------------------------------------------------------------------------

295 # Tests for submit_store

296 # ---------------------------------------------------------------------------

@@ -166,16 +298,16 @@ class TestLookup:

298

299 class TestSubmitStore:

300 def test_no_decode_succeeds_immediately(self):

169 - """Without a ``decode`` block, job succeeds immediately."""

301 + """Without a ``remote_decoder`` block, job succeeds immediately."""

302 mgr = _make_manager()

303 job = _job_metadata(job_id=1, kv_params={})

304 mgr.submit_store(job)

305 assert mgr._finished_jobs == [JobResult(job_id=1, success=True)]

306

307 def test_missing_kv_request_id_fails(self):

176 - """Missing kv_request_id inside ``decode`` fails the job."""

308 + """Missing kv_request_id inside ``remote_decoder`` fails the job."""

309 mgr = _make_manager()

178 - params: dict = {"decode": {}}

310 + params: dict = {"remote_decoder": {}}

311 job = _job_metadata(job_id=1, kv_params=params)

312 mgr.submit_store(job)

313 assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]

@@ -189,7 +321,7 @@ class TestSubmitStore:

321 job_id=1,

322 keys=[b"k1", b"k2"],

323 block_ids=[3, 4],

192 - kv_params=_decode_kv_params(kv_request_id="req-1"),

324 + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"),

325 )

326 mgr.submit_store(job)

327

@@ -220,7 +352,7 @@ class TestSubmitStore:

352 job_id=7,

353 keys=[b"k1", b"k2"],

354 block_ids=[3, 4],

223 - kv_params=_decode_kv_params(kv_request_id="req-1"),

355 + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"),

356 )

357 mgr.submit_store(job)

358

@@ -233,9 +365,9 @@ class TestSubmitStore:

365 def test_extra_top_level_keys_are_ignored(self):

366 """Producer-side kv_transfer_params should not pre-create a

367 session even when a stale caller still passes a top-level

236 - ``remote_host``/``remote_port`` next to ``decode``."""

368 + ``remote_host``/``remote_port`` next to ``remote_decoder``."""

369 mgr = _make_manager()

238 - params = _decode_kv_params()

370 + params = _remote_decoder_kv_params()

371 params["remote_host"] = "stale"

372 params["remote_port"] = 12345

373 job = _job_metadata(job_id=1, kv_params=params)

@@ -262,7 +394,7 @@ class TestSubmitLoad:

394 """Empty key list succeeds immediately."""

395 mgr = _make_manager()

396 job = _job_metadata(

265 - job_id=1, keys=[], block_ids=[], kv_params=_prefill_kv_params()

397 + job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params()

398 )

399 mgr.submit_load(job)

400 assert mgr._finished_jobs == [JobResult(job_id=1, success=True)]

@@ -270,7 +402,7 @@ class TestSubmitLoad:

402 def test_no_session_fails(self):

403 """No session for peer fails and marks request failed."""

404 mgr = _make_manager()

273 - job = _job_metadata(job_id=1, kv_params=_prefill_kv_params())

405 + job = _job_metadata(job_id=1, kv_params=_remote_prefiller_kv_params())

406 mgr.submit_load(job)

407 assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]

408 assert "req-1" in mgr._failed_req_ids

@@ -286,7 +418,7 @@ class TestSubmitLoad:

418 job_id=42,

419 keys=[b"k1", b"k2"],

420 block_ids=[5, 6],

289 - kv_params=_prefill_kv_params(kv_request_id="req-42"),

421 + kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"),

422 )

423 mgr.submit_load(job)

424

@@ -294,6 +426,20 @@ class TestSubmitLoad:

426 assert mgr._finished_jobs == []

427 assert "req-42" not in mgr._failed_req_ids

428

429 + def test_missing_consumer_flag_fails(self):

430 + """Peer fields present but neither do_remote_prefill nor

431 + do_p2p_fetch is set — submit_load fails the job rather than

432 + emit a stray FetchMsg."""

433 + mgr = _make_manager()

434 + params = {

435 + "remote_host": "10.0.0.1",

436 + "remote_port": 8000,

437 + "kv_request_id": "req-1",

438 + }

439 + job = _job_metadata(job_id=1, kv_params=params)

440 + mgr.submit_load(job)

441 + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]

442 +

443

444 # ---------------------------------------------------------------------------

445 # Tests for on_request_finished

@@ -308,7 +454,7 @@ class TestOnRequestFinished:

454

455 def test_prunes_failed_req_ids(self):

456 mgr = self._make_with_failed()

311 - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))

457 + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1"))

458 mgr.on_request_finished(ctx)

459 assert "req-1" not in mgr._failed_req_ids

460

@@ -325,14 +471,26 @@ class TestOnRequestFinished:

471 assert "req-1" in mgr._failed_req_ids

472

473 def test_decoder_side_calls_session_finish_request(self):

328 - """Decoder-side finish (``prefill`` set) still routes via peer_id

474 + """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id

475 because the consumer addresses the producer it loaded from. The

476 session's finish_request cancels the client-role load."""

477 mgr = _make_manager()

478 peer_id = "10.0.0.1:8000"

479 session = _FakeSession(peer_id=peer_id)

480 mgr._sessions[peer_id] = session

335 - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))

481 + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1"))

482 + mgr.on_request_finished(ctx)

483 + assert session.finishes == ["req-1"]

484 +

485 + def test_p2p_consumer_side_calls_session_finish_request(self):

486 + """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id

487 + so the session drops any pending lookups (cancel_lookups) and

488 + cancels any inbound load."""

489 + mgr = _make_manager()

490 + peer_id = "10.0.0.1:8000"

491 + session = _FakeSession(peer_id=peer_id)

492 + mgr._sessions[peer_id] = session

493 + ctx = _req_context(kv_params=_remote_kv_source_kv_params(kv_request_id="req-1"))

494 mgr.on_request_finished(ctx)

495 assert session.finishes == ["req-1"]

496

@@ -342,7 +500,7 @@ class TestOnRequestFinished:

500 mgr = _make_manager()

501 bound = _FakeSession(peer_id="some-peer:1", connected=True)

502 mgr._kv_to_session["req-1"] = bound # type: ignore[assignment]

345 - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1"))

503 + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1"))

504 mgr.on_request_finished(ctx)

505 assert bound.finishes == ["req-1"]

506 assert "req-1" not in mgr._kv_to_session

@@ -360,7 +518,7 @@ class TestOnRequestFinished:

518 _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]),

519 _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]),

520 ]

363 - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1"))

521 + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1"))

522 mgr.on_request_finished(ctx)

523 assert "req-1" in mgr._unbound_stores

524 assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11]

@@ -378,11 +536,19 @@ class _FakeServerHalf:

536 def __init__(self) -> None:

537 self._inflight: dict[int, object] = {}

538

539 + @property

540 + def has_inflight_transfers(self) -> bool:

541 + return bool(self._inflight)

542 +

543

544 class _FakeClientHalf:

545 def __init__(self) -> None:

546 self._inbound: dict[int, object] = {}

547

548 + @property

549 + def has_active_loads(self) -> bool:

550 + return bool(self._inbound)

551 +

552

553 class _FakeSession:

554 """Fake bidirectional session that returns canned poll() results."""

@@ -395,8 +561,10 @@ class _FakeSession:

561 loads: list[LoadResult] | None = None,

562 stores: list[StoreResult] | None = None,

563 new_fetch_ids: list[str] | None = None,

398 - close_loads: list[tuple[int, str]] | None = None,

564 + close_jobs: list[int] | None = None,

565 + close_req_ids: list[str] | None = None,

566 close_stores: list[int] | None = None,

567 + close_failed_serves: list[ReqContext] | None = None,

568 ) -> None:

569 self.peer_id = peer_id

570 self.alive = alive

@@ -405,18 +573,24 @@ class _FakeSession:

573 self._loads = loads or []

574 self._stores = stores or []

575 self._new_fetch_ids = new_fetch_ids or []

408 - self._close_loads = close_loads or []

576 + self._close_jobs = close_jobs or []

577 + self._close_req_ids = close_req_ids or []

578 self._close_stores = close_stores or []

579 + self._close_failed_serves = close_failed_serves or []

580 self.requests: list[tuple[int, str]] = []

581 self.stores_added: list[tuple[str, list, object, int]] = []

582 self.attached: list[object] = []

583 self.finishes: list[str] = []

584 # Mirror P2PSession._server._inflight (transfer_id → handle) and

415 - # P2PSession._client._inbound for the shutdown-drain and drain_jobs

416 - # paths. Tests populate _server._inflight when needed.

585 + # P2PSession._client.has_active_loads for the shutdown-drain and

586 + # drain_jobs paths. Tests populate _server._inflight when needed.

587 self._server = _FakeServerHalf()

588 self._client = _FakeClientHalf()

589

590 + @property

591 + def has_pending_work(self) -> bool:

592 + return self._client.has_active_loads or self._server.has_inflight_transfers

593 +

594 def poll(self):

595 result = SessionPollResult(

596 loads=self._loads,

@@ -442,7 +616,12 @@ class _FakeSession:

616 self.finishes.append(kv_request_id)

617

618 def close(self):

445 - return self._close_loads, self._close_stores

619 + return SessionCloseResult(

620 + failed_jobs=self._close_jobs,

621 + failed_req_ids=self._close_req_ids,

622 + failed_stores=self._close_stores,

623 + failed_serves=self._close_failed_serves,

624 + )

625

626

627 class TestGetFinished:

@@ -484,7 +663,8 @@ class TestGetFinished:

663 peer_id="dead:1234",

664 alive=False,

665 connected=True,

487 - close_loads=[(20, "req-load")],

666 + close_jobs=[20],

667 + close_req_ids=["req-load"],

668 close_stores=[10, 11],

669 )

670 mgr._sessions["dead:1234"] = dead # type: ignore[assignment]

@@ -498,6 +678,29 @@ class TestGetFinished:

678 assert "dead:1234" not in mgr._sessions

679 assert "req-load" in mgr._failed_req_ids

680

681 + def test_reap_fails_probes(self):

682 + """A reaped session's in-flight lookups land in _failed_req_ids so

683 + the consumer's lookup() returns MISS instead of RETRY forever."""

684 +

685 + class FakeData:

686 + def remove_remote_peer(self, pid):

687 + pass

688 +

689 + mgr = self._make()

690 + mgr._data = FakeData() # type: ignore[assignment]

691 + dead = _FakeSession(

692 + peer_id="dead:1234",

693 + alive=False,

694 + connected=True,

695 + close_req_ids=["req-probe-1", "req-probe-2"],

696 + )

697 + mgr._sessions["dead:1234"] = dead # type: ignore[assignment]

698 +

699 + list(mgr.get_finished_jobs())

700 + assert "dead:1234" not in mgr._sessions

701 + assert "req-probe-1" in mgr._failed_req_ids

702 + assert "req-probe-2" in mgr._failed_req_ids

703 +

704 def test_unbound_store_kept_within_timeout(self):

705 """Recently-parked unbound stores stay across a poll."""

706 mgr = self._make()

@@ -537,7 +740,7 @@ class TestGetFinished:

740 submitted_at stamp so the unbound-store sweep can age it out."""

741 mgr = _make_manager()

742 job = _job_metadata(

540 - job_id=1, kv_params=_decode_kv_params(kv_request_id="req-1")

743 + job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1")

744 )

745 before = time.monotonic()

746 mgr.submit_store(job)

@@ -881,21 +1084,21 @@ class TestBidirectionalManager:

1084 b_loads_kv = "req-BtoA-load" # B loads, A serves

1085

1086 a_decoder_params = {

884 - "prefill": {

1087 + "remote_prefiller": {

1088 "kv_request_id": a_loads_kv,

1089 "remote_host": "B",

1090 "remote_port": 2,

1091 },

1092 }

1093 b_decoder_params = {

891 - "prefill": {

1094 + "remote_prefiller": {

1095 "kv_request_id": b_loads_kv,

1096 "remote_host": "A",

1097 "remote_port": 1,

1098 },

1099 }

897 - a_prefiller_params = {"decode": {"kv_request_id": b_loads_kv}}

898 - b_prefiller_params = {"decode": {"kv_request_id": a_loads_kv}}

1100 + a_prefiller_params = {"remote_decoder": {"kv_request_id": b_loads_kv}}

1101 + b_prefiller_params = {"remote_decoder": {"kv_request_id": a_loads_kv}}

1102

1103 # 1. Both sides open client-role sessions toward the peer.

1104 mgr_a.on_new_request(_req_context(a_decoder_params))

@@ -1065,7 +1268,8 @@ class TestPollOnce:

1268 peer_id=peer_dead,

1269 alive=False,

1270 connected=True,

1068 - close_loads=[(33, "req-33")],

1271 + close_jobs=[33],

1272 + close_req_ids=["req-33"],

1273 close_stores=[44],

1274 )

1275 mgr._sessions[peer_dead] = dead # type: ignore[assignment]

@@ -1309,13 +1513,13 @@ class TestConnectionDeathMidTransfer:

1513 mgr_a, mgr_b = _build_paired_managers()

1514

1515 a_decoder_params = {

1312 - "prefill": {

1516 + "remote_prefiller": {

1517 "kv_request_id": "req-load",

1518 "remote_host": "B",

1519 "remote_port": 2,

1520 },

1521 }

1318 - a_prefiller_params = {"decode": {"kv_request_id": "req-store"}}

1522 + a_prefiller_params = {"remote_decoder": {"kv_request_id": "req-store"}}

1523

1524 # Open the outbound session A->B and submit one load + one store.

1525 mgr_a.on_new_request(_req_context(a_decoder_params))

@@ -1394,6 +1598,7 @@ class TestBindHostPortDefaults:

1598 identity (``host:port``) stays decoupled from the NIXL agent name

1599 (a uuid).

1600 """

1601 + monkeypatch.setenv("PYTHONHASHSEED", "0")

1602 monkeypatch.setattr(

1603 manager_module,

1604 "FileMapper",

tests/v1/kv_offload/tiering/p2p/test_sessions.py

+966 -83

@@ -13,9 +13,18 @@ completes its own load.

13 from __future__ import annotations

15 import time

16 +from collections.abc import Sequence

18 +import numpy as np

19 import pytest

21 +from vllm.v1.kv_offload.base import (

22 + LookupResult,

23 + OffloadKey,

24 + ReqContext,

25 + RequestOffloadingContext,

26 +)

27 +from vllm.v1.kv_offload.tiering.base import JobMetadata

28 from vllm.v1.kv_offload.tiering.p2p.session import (

29 LoadResult,

30 P2PSession,

@@ -33,6 +42,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import (

42 ConnectMsg,

43 DisconnectMsg,

44 FetchMsg,

45 + LookupMsg,

46 + LookupRespMsg,

47 TransferDoneMsg,

48 )

49 from vllm.v1.kv_offload.tiering.p2p.session.server import (

@@ -48,6 +59,11 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import (

59 # ---------------------------------------------------------------------------

62 +# Shared PYTHONHASHSEED used by the session under test and the fake peer's

63 +# ConnectMsg so the handshake succeeds unless a test overrides one side.

64 +_DEFAULT_HASH_SEED = "0"

65 +

66 +

67 class FakeDataTransport:

68 """Minimal fake DataTransport for testing sessions."""

@@ -146,12 +162,16 @@ class FakeConnection:

162 self._inbox: list[dict] = []

163 self._sent: list[dict] = []

164 self._closed = False

165 + # When True, send() raises to simulate a broken/dead connection.

166 + self.fail_send = False

167

168 @property

169 def alive(self) -> bool:

170 return not self._closed

171

172 def send(self, msg: dict) -> None:

173 + if self.fail_send:

174 + raise ConnectionError("simulated dead connection")

175 self._sent.append(msg)

176

177 def recv(self) -> list[dict]:

@@ -173,6 +193,7 @@ def _peer_connect_msg(

193 peer_id: str = "peer:8000",

194 block_len: int = 4096,

195 fingerprint: str | None = None,

196 + hash_seed: str = _DEFAULT_HASH_SEED,

197 ) -> dict:

198 """Build a ConnectMsg as if the peer sent it."""

199 msg = {

@@ -182,17 +203,82 @@ def _peer_connect_msg(

203 ConnectMsg.BASE_ADDR: 0x2000,

204 ConnectMsg.NUM_BLOCKS: 16,

205 ConnectMsg.BLOCK_LEN: block_len,

206 + ConnectMsg.HASH_SEED: hash_seed,

207 }

208 if fingerprint is not None:

209 msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint

210 return msg

211

212

213 +class FakeParent:

214 + """Configurable :class:`ParentManager` for server-role tests.

215 +

216 + ``stored`` is the dict of ready blocks (key → primary block_id).

217 + ``pending`` and ``retry`` script the first lookup() result for those

218 + keys; subsequent lookups behave normally (a key that promised

219 + HIT_PENDING / RETRY can later be promoted to HIT by adding it to

220 + ``stored`` and removing it from ``pending``/``retry``). ``calls``

221 + captures every parent invocation in order for assertions.

222 +

223 + Injected per-step via ``session.serve_external_requests(parent)`` —

224 + not held by the session, matching how ``TieringOffloadingManager``

225 + hands the tier a handle valid only for that call.

226 + """

227 +

228 + def __init__(

229 + self,

230 + stored: dict[OffloadKey, int] | None = None,

231 + pending: set[OffloadKey] | None = None,

232 + retry: set[OffloadKey] | None = None,

233 + ) -> None:

234 + self.stored: dict[OffloadKey, int] = dict(stored or {})

235 + self.pending: set[OffloadKey] = set(pending or ())

236 + self.retry: set[OffloadKey] = set(retry or ())

237 + self._next_job_id: int = 1000

238 + self.calls: list[tuple] = []

239 +

240 + def on_new_request(self, ctx: ReqContext) -> RequestOffloadingContext:

241 + self.calls.append(("on_new_request", ctx.req_id))

242 + return RequestOffloadingContext()

243 +

244 + def lookup(self, key: OffloadKey, ctx: ReqContext) -> LookupResult:

245 + self.calls.append(("lookup", key, ctx.req_id))

246 + if key in self.pending:

247 + return LookupResult.HIT_PENDING

248 + if key in self.retry:

249 + return LookupResult.RETRY

250 + if key in self.stored:

251 + return LookupResult.HIT

252 + return LookupResult.MISS

253 +

254 + def create_store_job(

255 + self,

256 + keys: Sequence[OffloadKey],

257 + ctx: ReqContext,

258 + ) -> JobMetadata:

259 + keys_list = list(keys)

260 + self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id))

261 + block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32)

262 + job_id = self._next_job_id

263 + self._next_job_id += 1

264 + return JobMetadata(

265 + job_id=job_id,

266 + keys=keys_list,

267 + block_ids=block_ids,

268 + is_promotion=False,

269 + req_context=ctx,

270 + )

271 +

272 + def on_request_finished(self, ctx: ReqContext) -> None:

273 + self.calls.append(("on_request_finished", ctx.req_id))

274 +

275 +

276 def _make_session(

277 conn: FakeConnection | None = None,

278 transport: FakeDataTransport | None = None,

279 peer_id: str = "peer:8000",

280 local_id: str = "local:9000",

281 + local_hash_seed: str = _DEFAULT_HASH_SEED,

282 ) -> tuple[P2PSession, FakeConnection, FakeDataTransport]:

283 if conn is None:

284 conn = FakeConnection(peer_id=peer_id)

@@ -203,11 +289,17 @@ def _make_session(

289 local_id=local_id,

290 transport=transport, # type: ignore[arg-type]

291 local_block_len=transport.block_len,

292 + local_hash_seed=local_hash_seed,

293 conn=conn, # type: ignore[arg-type]

294 )

295 return session, conn, transport

296

297

298 +def _serve(session: P2PSession, parent: FakeParent) -> None:

299 + """Resolve enqueued inbound lookups, as the manager does each step."""

300 + session.serve_external_requests(parent) # type: ignore[arg-type]

301 +

302 +

303 def _activate(

304 session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000"

305 ) -> None:

@@ -217,6 +309,42 @@ def _activate(

309 session.poll()

310

311

312 +# --- Accessors for the server role's consolidated per-request state.

313 +# ServerRole keeps one _ServerRequestState per kv_request_id; entries are

314 +# garbage-collected once fully idle, so "no outbound"/"no abort" reads as

315 +# either a missing entry or a None field. These helpers paper over that.

316 +

317 +

318 +def _srv_outbound(session: P2PSession, kv_request_id: str):

319 + """Outbound serve state for a kv_request_id, or None (idle / GC'd)."""

320 + st = session._server._requests.get(kv_request_id)

321 + return st.outbound if st is not None else None

322 +

323 +

324 +def _srv_lookups(session: P2PSession) -> list:

325 + """Every parked inbound _ActiveLookup across all requests."""

326 + return [

327 + lu for st in session._server._requests.values() for lu in st.lookups.values()

328 + ]

329 +

330 +

331 +def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None:

332 + """Pending-abort start time for a kv_request_id, or None."""

333 + st = session._server._requests.get(kv_request_id)

334 + return st.abort_started_at if st is not None else None

335 +

336 +

337 +def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int:

338 + """Inflight-transfer count tracked for a kv_request_id (0 if idle)."""

339 + st = session._server._requests.get(kv_request_id)

340 + return len(st.inflight_tids) if st is not None else 0

341 +

342 +

343 +def _srv_total_inflight(session: P2PSession) -> int:

344 + """Sum of per-request inflight counts across all requests."""

345 + return sum(len(st.inflight_tids) for st in session._server._requests.values())

346 +

347 +

348 # ---------------------------------------------------------------------------

349 # Connect / handshake

350 # ---------------------------------------------------------------------------

@@ -300,6 +428,29 @@ class TestConnectHandshake:

428 session.poll()

429 assert "peer:8000" in transport._remote_peers

430

431 + def test_hash_seed_mismatch_marks_dead(self):

432 + """Mismatched PYTHONHASHSEED rejects peer and marks connection dead."""

433 + session, conn, transport = _make_session(local_hash_seed="0")

434 + conn.enqueue(_peer_connect_msg(hash_seed="12345")) # mismatch

435 + session.poll()

436 + assert "peer:8000" not in transport._remote_peers

437 + assert not session.alive

438 + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent)

439 +

440 + def test_hash_seed_match_succeeds(self):

441 + """Matching PYTHONHASHSEED registers the peer and acks."""

442 + session, conn, transport = _make_session(local_hash_seed="12345")

443 + conn.enqueue(_peer_connect_msg(hash_seed="12345"))

444 + session.poll()

445 + assert "peer:8000" in transport._remote_peers

446 + assert session.alive

447 + assert any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent)

448 +

449 + def test_hash_seed_advertised_in_connect_msg(self):

450 + """Session advertises its own PYTHONHASHSEED in the ConnectMsg."""

451 + _, conn, _ = _make_session(local_hash_seed="777")

452 + assert conn._sent[0][ConnectMsg.HASH_SEED] == "777"

453 +

454

455 # ---------------------------------------------------------------------------

456 # Client-role flows

@@ -316,7 +467,7 @@ class TestClientFlows:

467 lookup = conn._sent[-1]

468 assert lookup[TYPE_KEY] == FetchMsg.TYPE

469 assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1"

319 - assert lookup[FetchMsg.BLOCK_HASHES] == [b"k1", b"k2"]

470 + assert lookup[FetchMsg.KEYS] == [b"k1", b"k2"]

471 assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1]

472

473 def test_transfer_done_success(self):

@@ -362,13 +513,46 @@ class TestClientFlows:

513 assert abort[TYPE_KEY] == AbortFetchMsg.TYPE

514 assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1"

515

516 + def test_active_loads_work_list_tracks_in_flight(self):

517 + """collect_results / has_active_loads use the _active_loads work-list,

518 + armed when a fetch is issued and discarded exactly when its load

519 + clears — a probe-only request never enters it, and completion empties

520 + it while the entry may briefly linger for GC."""

521 + session, conn, _ = _make_session()

522 + _activate(session, conn)

523 + client = session._client

524 +

525 + # A probe-only request has no in-flight load: not in _active_loads.

526 + session.register_lookup("req-probe", b"hp")

527 + assert client._active_loads == set()

528 + assert client.has_active_loads is False

529 +

530 + # Issuing a fetch arms the work-list.

531 + session.request_blocks(

532 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]

533 + )

534 + assert client._active_loads == {"req-1"}

535 + assert client.has_active_loads is True

536 +

537 + # Completion clears the load and discards it from the work-list.

538 + conn.enqueue(

539 + {

540 + TYPE_KEY: TransferDoneMsg.TYPE,

541 + TransferDoneMsg.KV_REQUEST_ID: "req-1",

542 + TransferDoneMsg.SUCCESS: True,

543 + }

544 + )

545 + session.poll()

546 + assert client._active_loads == set()

547 + assert client.has_active_loads is False

548 +

549 def test_load_timeout_sends_abort(self):

550 session, conn, _ = _make_session()

551 _activate(session, conn)

552 session.request_blocks(

553 job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]

554 )

371 - session._client._inbound["req-1"].submitted_at = time.monotonic() - 60.0

555 + session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0

556 session.poll()

557 abort = conn._sent[-1]

558 assert abort[TYPE_KEY] == AbortFetchMsg.TYPE

@@ -376,7 +560,7 @@ class TestClientFlows:

560 def test_load_abort_ack_timeout_surfaces_failure(self):

561 """After load timeout sends AbortFetch, if no AbortAck arrives within

562 _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

563 + from _requests — the engine cannot wait forever on a peer that won't

564 ack.

565 """

566 session, conn, _ = _make_session()

@@ -385,7 +569,7 @@ class TestClientFlows:

569 job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0]

570 )

571 # 1) Trip the load timeout to send AbortFetch and stamp aborted_at.

388 - session._client._inbound["req-7"].submitted_at = (

572 + session._client._requests["req-7"].load.submitted_at = (

573 time.monotonic() - _LOAD_TIMEOUT_S - 1.0

574 )

575 loads = session.poll().loads

@@ -395,32 +579,32 @@ class TestClientFlows:

579 and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7"

580 for m in conn._sent

581 )

398 - assert session._client._inbound["req-7"].aborted_at is not None

582 + assert session._client._requests["req-7"].load.aborted_at is not None

583

584 # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever

585 # arrived from the peer.

402 - session._client._inbound["req-7"].aborted_at = (

586 + session._client._requests["req-7"].load.aborted_at = (

587 time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0

588 )

589 loads = session.poll().loads

590 assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)]

407 - assert "req-7" not in session._client._inbound

591 + assert "req-7" not in session._client._requests

592

593 def test_load_abort_ack_clears_request(self):

594 """After load timeout sends AbortFetch, an arriving AbortAckMsg from

595 the peer surfaces the failure cleanly and removes the request from

412 - _inbound — covers the on_abort_ack arrival path."""

596 + _requests — covers the on_abort_ack arrival path."""

597 session, conn, _ = _make_session()

598 _activate(session, conn)

599 session.request_blocks(

600 job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0]

601 )

418 - session._client._inbound["req-8"].submitted_at = (

602 + session._client._requests["req-8"].load.submitted_at = (

603 time.monotonic() - _LOAD_TIMEOUT_S - 1.0

604 )

605 # First poll: AbortFetch goes out.

606 session.poll()

423 - assert session._client._inbound["req-8"].aborted_at is not None

607 + assert session._client._requests["req-8"].load.aborted_at is not None

608

609 # Peer acks the abort.

610 conn.enqueue(

@@ -431,7 +615,652 @@ class TestClientFlows:

615 )

616 loads = session.poll().loads

617 assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)]

434 - assert "req-8" not in session._client._inbound

618 + assert "req-8" not in session._client._requests

619 +

620 +

621 +# ---------------------------------------------------------------------------

622 +# Symmetric-P2P lookup flow (do_p2p_fetch)

623 +# ---------------------------------------------------------------------------

624 +

625 +

626 +class TestLookupFlow:

627 + """Consumer-side state machine for do_p2p_fetch lookups."""

628 +

629 + def test_aggregate_flush_resolve_round_trip(self):

630 + """register_lookup → flush sends one LookupMsg → response

631 + resolves entries → register_lookup returns the cached bool on

632 + every call, and repeat probes never re-issue a LookupMsg."""

633 + session, conn, _ = _make_session()

634 + _activate(session, conn)

635 +

636 + # Aggregate two keys for the same kv_request_id; both return None.

637 + assert session.register_lookup("req-1", b"hA") is None

638 + assert session.register_lookup("req-1", b"hB") is None

639 +

640 + # Flush sends one LookupMsg with both keys.

641 + sent_before = len(conn._sent)

642 + session.flush_pending_lookups()

643 + new = conn._sent[sent_before:]

644 + assert len(new) == 1

645 + msg = new[0]

646 + assert msg[TYPE_KEY] == LookupMsg.TYPE

647 + assert msg[LookupMsg.KV_REQUEST_ID] == "req-1"

648 + assert sorted(msg[LookupMsg.KEYS]) == [b"hA", b"hB"]

649 +

650 + # Idempotent re-flush sends nothing — the entries are now in-flight.

651 + sent_before = len(conn._sent)

652 + session.flush_pending_lookups()

653 + assert conn._sent[sent_before:] == []

654 +

655 + # While in-flight, register_lookup keeps returning None.

656 + assert session.register_lookup("req-1", b"hA") is None

657 +

658 + # Peer answers: hA hit, hB miss.

659 + conn.enqueue(

660 + {

661 + TYPE_KEY: LookupRespMsg.TYPE,

662 + LookupRespMsg.KV_REQUEST_ID: "req-1",

663 + LookupRespMsg.KEYS: [b"hA", b"hB"],

664 + LookupRespMsg.HITS: [True, False],

665 + }

666 + )

667 + session.poll()

668 +

669 + # register_lookup returns the resolved bool.

670 + assert session.register_lookup("req-1", b"hA") is True

671 + assert session.register_lookup("req-1", b"hB") is False

672 + # The entry is cached, not popped: repeat probes keep returning the

673 + # same result and never re-queue the key, so a flush sends nothing.

674 + assert session.register_lookup("req-1", b"hA") is True

675 + assert session.register_lookup("req-1", b"hB") is False

676 + sent_before = len(conn._sent)

677 + session.flush_pending_lookups()

678 + assert [

679 + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE

680 + ] == []

681 +

682 + def test_request_blocks_clears_probe_cache(self):

683 + """A resolved HIT probe is popped when its fetch is issued, so a

684 + re-scheduled request re-probes instead of trusting the stale True

685 + (the served block is unpinned and may have been evicted)."""

686 + session, conn, _ = _make_session()

687 + _activate(session, conn)

688 +

689 + # Probe hA, flush, and let the peer resolve it to a HIT.

690 + assert session.register_lookup("req-1", b"hA") is None

691 + session.flush_pending_lookups()

692 + conn.enqueue(

693 + {

694 + TYPE_KEY: LookupRespMsg.TYPE,

695 + LookupRespMsg.KV_REQUEST_ID: "req-1",

696 + LookupRespMsg.KEYS: [b"hA"],

697 + LookupRespMsg.HITS: [True],

698 + }

699 + )

700 + session.poll()

701 + assert session.register_lookup("req-1", b"hA") is True

702 +

703 + # Fetch consumes the probe.

704 + session.request_blocks(

705 + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[0]

706 + )

707 + assert b"hA" not in session._client._requests["req-1"].probes

708 +

709 + # Re-scheduled probe of the same key is treated as brand-new: it

710 + # returns None and re-queues, so a flush emits a fresh LookupMsg.

711 + assert session.register_lookup("req-1", b"hA") is None

712 + sent_before = len(conn._sent)

713 + session.flush_pending_lookups()

714 + fresh = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE]

715 + assert len(fresh) == 1

716 + assert fresh[0][LookupMsg.KEYS] == [b"hA"]

717 +

718 + def test_flush_uses_work_list_not_full_scan(self):

719 + """flush drains a work-list rather than scanning every live request.

720 +

721 + A request with no newly-registered keys is not revisited: after a

722 + flush the work-list is empty, an idle re-flush sends nothing, and a

723 + subsequent register re-arms exactly the one affected id — even while

724 + an unrelated request stays live in ``_requests``.

725 + """

726 + session, conn, _ = _make_session()

727 + _activate(session, conn)

728 + client = session._client

729 +

730 + # Two requests register keys; both are queued for flush.

731 + session.register_lookup("req-A", b"hA")

732 + session.register_lookup("req-B", b"hB")

733 + assert client._flush_pending == {"req-A", "req-B"}

734 +

735 + # Flush drains the work-list even though both requests stay live.

736 + session.flush_pending_lookups()

737 + assert client._flush_pending == set()

738 + assert set(client._requests) == {"req-A", "req-B"}

739 +

740 + # An idle re-flush visits nothing and sends no LookupMsg.

741 + sent_before = len(conn._sent)

742 + session.flush_pending_lookups()

743 + assert [

744 + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE

745 + ] == []

746 +

747 + # A new register re-arms only that id.

748 + session.register_lookup("req-A", b"hA2")

749 + assert client._flush_pending == {"req-A"}

750 +

751 + def test_separate_lookup_msg_per_kv_request_id(self):

752 + """Hashes for different kv_request_ids flush as separate LookupMsgs."""

753 + session, conn, _ = _make_session()

754 + _activate(session, conn)

755 +

756 + session.register_lookup("req-A", b"h1")

757 + session.register_lookup("req-B", b"h2")

758 + session.register_lookup("req-A", b"h3")

759 +

760 + sent_before = len(conn._sent)

761 + session.flush_pending_lookups()

762 + sent = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE]

763 + assert len(sent) == 2

764 + by_req = {m[LookupMsg.KV_REQUEST_ID]: m[LookupMsg.KEYS] for m in sent}

765 + assert sorted(by_req["req-A"]) == [b"h1", b"h3"]

766 + assert by_req["req-B"] == [b"h2"]

767 +

768 + def test_multiple_lookup_msgs_across_steps(self):

769 + """A request's block set may be discovered across scheduler steps:

770 + each step that registers new keys flushes its own LookupMsg for

771 + the same kv_request_id, carrying only the newly-probed keys."""

772 + session, conn, _ = _make_session()

773 + _activate(session, conn)

774 +

775 + # Step 1: probe hA, hB.

776 + session.register_lookup("req-1", b"hA")

777 + session.register_lookup("req-1", b"hB")

778 + sent_before = len(conn._sent)

779 + session.flush_pending_lookups()

780 + first = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE]

781 + assert len(first) == 1

782 + assert first[0][LookupMsg.KV_REQUEST_ID] == "req-1"

783 + assert sorted(first[0][LookupMsg.KEYS]) == [b"hA", b"hB"]

784 +

785 + # Step 2: a new key is discovered for the same request. The

786 + # in-flight keys from step 1 are not re-sent; a second LookupMsg

787 + # goes out carrying only the newly-probed key.

788 + assert session.register_lookup("req-1", b"hA") is None # in-flight no-op

789 + session.register_lookup("req-1", b"hC")

790 + sent_before = len(conn._sent)

791 + session.flush_pending_lookups()

792 + second = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE]

793 + assert len(second) == 1

794 + assert second[0][LookupMsg.KV_REQUEST_ID] == "req-1"

795 + assert second[0][LookupMsg.KEYS] == [b"hC"]

796 +

797 + def test_split_response_resolves_across_messages(self):

798 + """Producer may answer one LookupMsg's keys across multiple

799 + LookupRespMsgs — pairs are self-describing so each lands."""

800 + session, conn, _ = _make_session()

801 + _activate(session, conn)

802 +

803 + session.register_lookup("req-1", b"hA")

804 + session.register_lookup("req-1", b"hB")

805 + session.flush_pending_lookups()

806 +

807 + # Two responses, each carrying one of the two keys.

808 + conn.enqueue(

809 + {

810 + TYPE_KEY: LookupRespMsg.TYPE,

811 + LookupRespMsg.KV_REQUEST_ID: "req-1",

812 + LookupRespMsg.KEYS: [b"hA"],

813 + LookupRespMsg.HITS: [True],

814 + }

815 + )

816 + conn.enqueue(

817 + {

818 + TYPE_KEY: LookupRespMsg.TYPE,

819 + LookupRespMsg.KV_REQUEST_ID: "req-1",

820 + LookupRespMsg.KEYS: [b"hB"],

821 + LookupRespMsg.HITS: [False],

822 + }

823 + )

824 + session.poll()

825 +

826 + assert session.register_lookup("req-1", b"hA") is True

827 + assert session.register_lookup("req-1", b"hB") is False

828 +

829 + def test_finish_request_cancels_pending_lookups(self):

830 + """finish_request drops every pending lookup for the kv_request_id."""

831 + session, conn, _ = _make_session()

832 + _activate(session, conn)

833 +

834 + session.register_lookup("req-1", b"hA")

835 + session.register_lookup("req-1", b"hB")

836 + session.register_lookup("req-2", b"hC")

837 + session.finish_request("req-1")

838 +

839 + # req-1 entries gone, req-2 untouched.

840 + assert "req-1" not in session._client._requests

841 + assert b"hC" in session._client._requests["req-2"].probes

842 +

843 + def test_finish_after_flushed_lookup_sends_empty_fetch(self):

844 + """LookupMsg flushed but no FetchMsg sent (all-miss case) →

845 + finish_request emits an empty FetchMsg so the peer can drop its

846 + lookup state and call parent.on_request_finished."""

847 + session, conn, _ = _make_session()

848 + _activate(session, conn)

849 +

850 + session.register_lookup("req-1", b"hA")

851 + session.flush_pending_lookups()

852 + sent_before = len(conn._sent)

853 +

854 + session.finish_request("req-1")

855 +

856 + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE]

857 + assert len(fetches) == 1

858 + assert fetches[0][FetchMsg.KV_REQUEST_ID] == "req-1"

859 + assert fetches[0][FetchMsg.KEYS] == []

860 + assert fetches[0][FetchMsg.BLOCK_INDEXES] == []

861 +

862 + def test_finish_without_flushed_lookup_sends_no_fetch(self):

863 + """No LookupMsg was ever sent → finish_request must not emit an

864 + empty FetchMsg (the peer has no state to release)."""

865 + session, conn, _ = _make_session()

866 + _activate(session, conn)

867 +

868 + # Register but never flush.

869 + session.register_lookup("req-1", b"hA")

870 + sent_before = len(conn._sent)

871 +

872 + session.finish_request("req-1")

873 +

874 + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE]

875 + assert fetches == []

876 +

877 + def test_finish_after_real_fetch_sends_no_second_fetch(self):

878 + """A real FetchMsg was already sent for the id → finish_request

879 + must not emit a second (empty) FetchMsg."""

880 + session, conn, _ = _make_session()

881 + _activate(session, conn)

882 +

883 + session.register_lookup("req-1", b"hA")

884 + session.flush_pending_lookups()

885 + # Resolve the probe to a HIT before fetching, as the manager only

886 + # loads confirmed hits (an unresolved probe yields RETRY).

887 + conn.enqueue(

888 + {

889 + TYPE_KEY: LookupRespMsg.TYPE,

890 + LookupRespMsg.KV_REQUEST_ID: "req-1",

891 + LookupRespMsg.KEYS: [b"hA"],

892 + LookupRespMsg.HITS: [True],

893 + }

894 + )

895 + session.poll()

896 + session.request_blocks(

897 + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[7]

898 + )

899 + sent_before = len(conn._sent)

900 +

901 + session.finish_request("req-1")

902 +

903 + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE]

904 + assert fetches == []

905 +

906 + def test_server_lookup_deferred_until_serve_then_all_misses(self):

907 + """``poll()`` only enqueues an inbound LookupMsg — no response is

908 + sent until ``serve_external_requests``. With an all-miss parent

909 + the aggregated LookupRespMsg carries the same keys and

910 + ``hits=[False, ...]``."""

911 + session, conn, _ = _make_session()

912 + _activate(session, conn)

913 +

914 + sent_before = len(conn._sent)

915 + conn.enqueue(

916 + {

917 + TYPE_KEY: LookupMsg.TYPE,

918 + LookupMsg.KV_REQUEST_ID: "req-1",

919 + LookupMsg.KEYS: [b"hX", b"hY", b"hZ"],

920 + }

921 + )

922 + session.poll()

923 +

924 + # Dispatch alone must not answer — the parent handle is only valid

925 + # during serve_external_requests.

926 + assert [

927 + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE

928 + ] == []

929 +

930 + _serve(session, FakeParent())

931 +

932 + resps = [

933 + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE

934 + ]

935 + assert len(resps) == 1

936 + resp = resps[0]

937 + assert resp[LookupRespMsg.KV_REQUEST_ID] == "req-1"

938 + assert resp[LookupRespMsg.KEYS] == [b"hX", b"hY", b"hZ"]

939 + assert resp[LookupRespMsg.HITS] == [False, False, False]

940 +

941 +

942 +# ---------------------------------------------------------------------------

943 +# Server-side handling of inbound LookupMsg (ParentManager-driven)

944 +#

945 +# poll() only enqueues the LookupMsg; serve_external_requests(parent)

946 +# resolves it. Tests follow the poll() → _serve() pattern.

947 +# ---------------------------------------------------------------------------

948 +

949 +

950 +def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]):

951 + conn.enqueue(

952 + {

953 + TYPE_KEY: LookupMsg.TYPE,

954 + LookupMsg.KV_REQUEST_ID: kv_request_id,

955 + LookupMsg.KEYS: list(keys),

956 + }

957 + )

958 +

959 +

960 +def _lookup_resps(conn: FakeConnection, since: int = 0) -> list[dict]:

961 + return [m for m in conn._sent[since:] if m[TYPE_KEY] == LookupRespMsg.TYPE]

962 +

963 +

964 +class TestServerLookupHandling:

965 + def test_immediate_hits_create_one_store_job(self):

966 + """All-HIT batch: one create_store_job call with all keys, one

967 + LookupRespMsg with hits=[True]*N, on_request_finished fires at the

968 + end of serve, and `available` is populated for the eventual fetch."""

969 + cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3})

970 + session, conn, _ = _make_session()

971 + _activate(session, conn)

972 +

973 + sent_before = len(conn._sent)

974 + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC"])

975 + session.poll()

976 + _serve(session, cb)

977 +

978 + resps = _lookup_resps(conn, sent_before)

979 + assert len(resps) == 1

980 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB", b"hC"]

981 + assert resps[0][LookupRespMsg.HITS] == [True, True, True]

982 +

983 + kinds = [c[0] for c in cb.calls]

984 + assert kinds.count("create_store_job") == 1

985 + cs = next(c for c in cb.calls if c[0] == "create_store_job")

986 + assert cs[1] == (b"hA", b"hB", b"hC")

987 + assert cb.calls[-1][0] == "on_request_finished"

988 +

989 + # Hits are pinned in outbound state for the upcoming FetchMsg match.

990 + assert set(_srv_outbound(session, "req-1").available) == {

991 + b"hA",

992 + b"hB",

993 + b"hC",

994 + }

995 +

996 + def test_all_misses_no_store_job_finish_fires(self):

997 + """All-MISS batch: no create_store_job call; one LookupRespMsg

998 + with hits=[False]*N; on_request_finished fires at end of serve."""

999 + cb = FakeParent()

1000 + session, conn, _ = _make_session()

1001 + _activate(session, conn)

1002 +

1003 + sent_before = len(conn._sent)

1004 + _send_lookup(conn, "req-1", [b"hA", b"hB"])

1005 + session.poll()

1006 + _serve(session, cb)

1007 +

1008 + resps = _lookup_resps(conn, sent_before)

1009 + assert len(resps) == 1

1010 + assert resps[0][LookupRespMsg.HITS] == [False, False]

1011 + assert all(c[0] != "create_store_job" for c in cb.calls)

1012 + assert cb.calls[-1][0] == "on_request_finished"

1013 +

1014 + def test_mixed_hit_miss_pending_defers_response_until_aggregate(self):

1015 + """HIT/MISS resolutions do not go out on first sight when any

1016 + key is still HIT_PENDING / RETRY. The lookup parks until every

1017 + key has settled (or the deadline fires), then one

1018 + LookupRespMsg carries all keys in wire order."""

1019 + cb = FakeParent(

1020 + stored={b"hA": 1},

1021 + pending={b"hB"},

1022 + retry={b"hD"},

1023 + )

1024 + session, conn, _ = _make_session()

1025 + _activate(session, conn)

1026 +

1027 + sent_before = len(conn._sent)

1028 + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC", b"hD"])

1029 + session.poll()

1030 + _serve(session, cb)

1031 +

1032 + # No LookupRespMsg yet — hB and hD are still pending.

1033 + assert _lookup_resps(conn, sent_before) == []

1034 + # HIT is still pinned immediately so the eventual FetchMsg matches.

1035 + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"]

1036 + assert len(cs_calls) == 1

1037 + assert cs_calls[0][1] == (b"hA",)

1038 + # Lookup is parked; on_request_finished not yet called.

1039 + assert all(c[0] != "on_request_finished" for c in cb.calls)

1040 + assert len(_srv_lookups(session)) == 1

1041 +

1042 + def test_pending_resolves_then_aggregate_response_fires(self):

1043 + """A HIT_PENDING key that becomes HIT on a later poll releases

1044 + the deferred aggregate response: one LookupRespMsg carrying

1045 + both keys in wire order, and one create_store_job call per

1046 + HIT (the second HIT is pinned when it resolves, not when the

1047 + response goes out)."""

1048 + cb = FakeParent(stored={b"hA": 1}, pending={b"hB"})

1049 + session, conn, _ = _make_session()

1050 + _activate(session, conn)

1051 +

1052 + sent_before = len(conn._sent)

1053 + _send_lookup(conn, "req-1", [b"hA", b"hB"])

1054 + session.poll()

1055 + _serve(session, cb)

1056 + # No response yet — hB still pending.

1057 + assert _lookup_resps(conn, sent_before) == []

1058 +

1059 + # Promote hB.

1060 + cb.pending.discard(b"hB")

1061 + cb.stored[b"hB"] = 2

1062 +

1063 + # Drive resolver via a second serve_external_requests.

1064 + _serve(session, cb)

1065 +

1066 + resps = _lookup_resps(conn, sent_before)

1067 + assert len(resps) == 1

1068 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB"]

1069 + assert resps[0][LookupRespMsg.HITS] == [True, True]

1070 +

1071 + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"]

1072 + assert len(cs_calls) == 2

1073 + assert cs_calls[0][1] == (b"hA",)

1074 + assert cs_calls[1][1] == (b"hB",)

1075 + # on_request_finished fires once after the aggregate resolve.

1076 + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1

1077 + assert b"hA" in _srv_outbound(session, "req-1").available

1078 + assert b"hB" in _srv_outbound(session, "req-1").available

1079 +

1080 + def test_pending_timeout_replies_miss_no_store_job(self):

1081 + """A HIT_PENDING key that stays pending past the batch

1082 + ``deadline`` is force-MISS and never pinned; the deferred

1083 + aggregate response fires with hits=[False]."""

1084 + cb = FakeParent(pending={b"hA"})

1085 + session, conn, _ = _make_session()

1086 + _activate(session, conn)

1087 +

1088 + sent_before = len(conn._sent)

1089 + _send_lookup(conn, "req-1", [b"hA"])

1090 + session.poll()

1091 + _serve(session, cb)

1092 + # Initial serve: nothing immediate, lookup parked, no LookupRespMsg.

1093 + assert _lookup_resps(conn, sent_before) == []

1094 +

1095 + # Forge the deadline into the past to trigger the timeout branch.

1096 + lookup = _srv_lookups(session)[0]

1097 + lookup.deadline = time.monotonic() - 0.1

1098 +

1099 + _serve(session, cb)

1100 +

1101 + resps = _lookup_resps(conn, sent_before)

1102 + assert len(resps) == 1

1103 + assert resps[0][LookupRespMsg.KEYS] == [b"hA"]

1104 + assert resps[0][LookupRespMsg.HITS] == [False]

1105 + assert all(c[0] != "create_store_job" for c in cb.calls)

1106 + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1

1107 +

1108 + def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self):

1109 + """Two LookupMsgs for the same kv_request_id get distinct ctxs

1110 + and two on_request_finished calls (one per batch)."""

1111 + cb = FakeParent(stored={b"hA": 1, b"hB": 2})

1112 + session, conn, _ = _make_session()

1113 + _activate(session, conn)

1114 +

1115 + _send_lookup(conn, "req-1", [b"hA"])

1116 + session.poll()

1117 + _serve(session, cb)

1118 + _send_lookup(conn, "req-1", [b"hB"])

1119 + session.poll()

1120 + _serve(session, cb)

1121 +

1122 + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"]

1123 + assert len(finish_calls) == 2

1124 + # Distinct synthetic req_ids

1125 + assert finish_calls[0][1] != finish_calls[1][1]

1126 + # Both namespaced under the same kv_request_id

1127 + assert ":req-1:" in finish_calls[0][1]

1128 + assert ":req-1:" in finish_calls[1][1]

1129 +

1130 + def test_close_returns_open_batch_ctxs_as_failed_serves(self):

1131 + """Tearing the session down with a parked batch returns the

1132 + synthetic ctx as a failed serve (no parent handle at teardown) so

1133 + the manager can release the TieringManager's state on its next

1134 + serve."""

1135 + cb = FakeParent(pending={b"hA"})

1136 + session, conn, _ = _make_session()

1137 + _activate(session, conn)

1138 +

1139 + _send_lookup(conn, "req-1", [b"hA"])

1140 + session.poll()

1141 + _serve(session, cb)

1142 + assert len(_srv_lookups(session)) == 1

1143 + assert all(c[0] != "on_request_finished" for c in cb.calls)

1144 +

1145 + result = session.close()

1146 +

1147 + assert len(result.failed_serves) == 1

1148 + assert ":req-1:" in result.failed_serves[0].req_id

1149 + # close() itself must not call the parent.

1150 + assert all(c[0] != "on_request_finished" for c in cb.calls)

1151 +

1152 + def test_wire_finish_drops_pending_batches_for_kv_request_id(self):

1153 + """``ServerRole.finish(kv_request_id)`` drops every parked batch

1154 + whose kv_request_id matches and queues its ctx for the next

1155 + serve's on_request_finished."""

1156 + cb = FakeParent(pending={b"hA", b"hB"})

1157 + session, conn, _ = _make_session()

1158 + _activate(session, conn)

1159 +

1160 + _send_lookup(conn, "req-1", [b"hA"])

1161 + session.poll()

1162 + _serve(session, cb)

1163 + _send_lookup(conn, "req-2", [b"hB"])

1164 + session.poll()

1165 + _serve(session, cb)

1166 + assert len(_srv_lookups(session)) == 2

1167 +

1168 + session._server.finish("req-1")

1169 +

1170 + # req-1 batch dropped from parked lookups; its ctx queued for release.

1171 + remaining_kv_request_ids = {b.kv_request_id for b in _srv_lookups(session)}

1172 + assert remaining_kv_request_ids == {"req-2"}

1173 + queued = session._server._finished_lookup_ctxs

1174 + assert len(queued) == 1

1175 + assert ":req-1:" in queued[0].req_id

1176 +

1177 + # The next serve fires on_request_finished exactly once for req-1.

1178 + _serve(session, cb)

1179 + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"]

1180 + assert len(finish_calls) == 1

1181 + assert ":req-1:" in finish_calls[0][1]

1182 +

1183 + def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self):

1184 + """A peer FetchMsg terminates the lookup phase for its id: parked

1185 + lookups with matching kv_request_id are dropped and their

1186 + ctx queued for on_request_finished; other kv_request_ids untouched."""

1187 + cb = FakeParent(pending={b"hA", b"hB"})

1188 + session, conn, _ = _make_session()

1189 + _activate(session, conn)

1190 +

1191 + _send_lookup(conn, "req-1", [b"hA"])

1192 + session.poll()

1193 + _serve(session, cb)

1194 + _send_lookup(conn, "req-2", [b"hB"])

1195 + session.poll()

1196 + _serve(session, cb)

1197 + assert len(_srv_lookups(session)) == 2

1198 +

1199 + # Empty FetchMsg: peer signals "lookup phase done" without asking

1200 + # for any blocks (the all-miss case).

1201 + conn.enqueue(

1202 + {

1203 + TYPE_KEY: FetchMsg.TYPE,

1204 + FetchMsg.KV_REQUEST_ID: "req-1",

1205 + FetchMsg.KEYS: [],

1206 + FetchMsg.BLOCK_INDEXES: [],

1207 + }

1208 + )

1209 + session.poll()

1210 +

1211 + remaining_kv_request_ids = {lu.kv_request_id for lu in _srv_lookups(session)}

1212 + assert remaining_kv_request_ids == {"req-2"}

1213 + # Dispatch queues the ctx but does not call the parent yet.

1214 + queued = session._server._finished_lookup_ctxs

1215 + assert len(queued) == 1

1216 + assert ":req-1:" in queued[0].req_id

1217 +

1218 + # The next serve fires on_request_finished exactly once for req-1.

1219 + _serve(session, cb)

1220 + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"]

1221 + assert len(finish_calls) == 1

1222 + assert ":req-1:" in finish_calls[0][1]

1223 +

1224 + def test_lookup_then_fetch_round_trip_emits_store_result(self):

1225 + """End-to-end: lookup pins primary slots → fetch matches them →

1226 + NIXL transfer completes → StoreResult surfaces with the

1227 + create_store_job's job_id (the engine releases the pin)."""

1228 + cb = FakeParent(stored={b"hA": 7, b"hB": 8})

1229 + session, conn, transport = _make_session()

1230 + _activate(session, conn)

1231 +

1232 + _send_lookup(conn, "req-1", [b"hA", b"hB"])

1233 + session.poll()

1234 + _serve(session, cb)

1235 + cs = next(c for c in cb.calls if c[0] == "create_store_job")

1236 + # FakeParent issues monotonic job_ids starting at 1000.

1237 + expected_job_id = 1000

1238 +

1239 + # Consumer issues FetchMsg on the resolved hits.

1240 + conn.enqueue(

1241 + {

1242 + TYPE_KEY: FetchMsg.TYPE,

1243 + FetchMsg.KV_REQUEST_ID: "req-1",

1244 + FetchMsg.KEYS: [b"hA", b"hB"],

1245 + FetchMsg.BLOCK_INDEXES: [20, 21],

1246 + }

1247 + )

1248 + session.poll()

1249 +

1250 + # NIXL write_blocks called with our pinned local block_ids.

1251 + assert len(transport._transfers) == 1

1252 + _, (_peer, local, remote) = next(iter(transport._transfers.items()))

1253 + assert local == [7, 8]

1254 + assert remote == [20, 21]

1255 +

1256 + # Drive the transport completion.

1257 + transport._poll_done.append(0)

1258 + result = session.poll()

1259 +

1260 + store_results = [s for s in result.stores if s.success]

1261 + assert any(s.job_id == expected_job_id for s in store_results)

1262 + # Sanity: kv mention in synthetic ctx.

1263 + assert cs[2].startswith("p2p:")

1264

1265

1266 # ---------------------------------------------------------------------------

@@ -449,7 +1278,7 @@ class TestServerFlows:

1278 {

1279 TYPE_KEY: FetchMsg.TYPE,

1280 FetchMsg.KV_REQUEST_ID: "req-1",

452 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

1281 + FetchMsg.KEYS: [b"k1", b"k2"],

1282 FetchMsg.BLOCK_INDEXES: [10, 11],

1283 }

1284 )

@@ -467,7 +1296,7 @@ class TestServerFlows:

1296 {

1297 TYPE_KEY: FetchMsg.TYPE,

1298 FetchMsg.KV_REQUEST_ID: "req-1",

470 - FetchMsg.BLOCK_HASHES: [b"k1"],

1299 + FetchMsg.KEYS: [b"k1"],

1300 FetchMsg.BLOCK_INDEXES: [5],

1301 }

1302 )

@@ -485,7 +1314,7 @@ class TestServerFlows:

1314 {

1315 TYPE_KEY: FetchMsg.TYPE,

1316 FetchMsg.KV_REQUEST_ID: "req-1",

488 - FetchMsg.BLOCK_HASHES: [b"k1"],

1317 + FetchMsg.KEYS: [b"k1"],

1318 FetchMsg.BLOCK_INDEXES: [5],

1319 }

1320 )

@@ -508,11 +1337,11 @@ class TestServerFlows:

1337 session.poll()

1338 ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)

1339 assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"

511 - assert "req-1" not in session._server._pending_aborts

1340 + assert _srv_abort_started(session, "req-1") is None

1341

1342 def test_abort_fetch_defers_ack_when_cancel_pending(self):

1343 """If cancel(mode='wait') reports still-inflight tids, the ack is

515 - deferred and the abort is parked in _pending_aborts."""

1344 + deferred and the abort is parked (abort_started_at set)."""

1345 session, conn, transport = _make_session()

1346 _activate(session, conn)

1347 # Seed an inflight transfer for req-1 that the transport pretends

@@ -533,7 +1362,7 @@ class TestServerFlows:

1362 session.poll()

1363

1364 assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)

536 - assert "req-1" in session._server._pending_aborts

1365 + assert _srv_abort_started(session, "req-1") is not None

1366 # First attempt happens inside _on_abort_fetch; the per-tick

1367 # drain runs again at the end of poll() — both are wait-mode.

1368 assert all(mode == "wait" for _, mode in transport._cancel_calls)

@@ -558,7 +1387,7 @@ class TestServerFlows:

1387 }

1388 )

1389 session.poll()

561 - assert "req-1" in session._server._pending_aborts

1390 + assert _srv_abort_started(session, "req-1") is not None

1391

1392 # Backend finishes draining: transport.poll() will return tid as

1393 # DONE, and the next cancel(mode='wait') call sees it's gone.

@@ -569,7 +1398,7 @@ class TestServerFlows:

1398

1399 ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)

1400 assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"

572 - assert "req-1" not in session._server._pending_aborts

1401 + assert _srv_abort_started(session, "req-1") is None

1402 assert tid not in session._server._inflight

1403

1404 def test_abort_fetch_force_cancels_after_timeout(self):

@@ -591,9 +1420,9 @@ class TestServerFlows:

1420 }

1421 )

1422 session.poll()

594 - assert "req-1" in session._server._pending_aborts

1423 + assert _srv_abort_started(session, "req-1") is not None

1424 # Backdate past the drain deadline.

596 - session._server._pending_aborts["req-1"] = (

1425 + session._server._requests["req-1"].abort_started_at = (

1426 time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0

1427 )

1428 # Even if the transport still claims it can't cancel, the

@@ -604,7 +1433,7 @@ class TestServerFlows:

1433

1434 ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)

1435 assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"

607 - assert "req-1" not in session._server._pending_aborts

1436 + assert _srv_abort_started(session, "req-1") is None

1437 assert tid not in session._server._inflight

1438 assert ([tid], "immediate") in transport._cancel_calls

1439

@@ -627,7 +1456,7 @@ class TestServerFlows:

1456 }

1457 )

1458 session.poll()

630 - first_started_at = session._server._pending_aborts["req-1"]

1459 + first_started_at = _srv_abort_started(session, "req-1")

1460

1461 # Second AbortFetchMsg for the same kv_request_id while still

1462 # draining must not reset the deadline.

@@ -638,7 +1467,7 @@ class TestServerFlows:

1467 }

1468 )

1469 session.poll()

641 - assert session._server._pending_aborts["req-1"] == first_started_at

1470 + assert _srv_abort_started(session, "req-1") == first_started_at

1471

1472 # Now let the drain succeed and confirm exactly one ack ever.

1473 transport._cancel_still_inflight.discard(tid)

@@ -669,7 +1498,7 @@ class TestServerFlows:

1498 {

1499 TYPE_KEY: FetchMsg.TYPE,

1500 FetchMsg.KV_REQUEST_ID: "req-1",

672 - FetchMsg.BLOCK_HASHES: [b"k1"],

1501 + FetchMsg.KEYS: [b"k1"],

1502 FetchMsg.BLOCK_INDEXES: [5],

1503 }

1504 )

@@ -701,7 +1530,7 @@ class TestServerFlows:

1530 {

1531 TYPE_KEY: FetchMsg.TYPE,

1532 FetchMsg.KV_REQUEST_ID: "req-1",

704 - FetchMsg.BLOCK_HASHES: [b"k1"],

1533 + FetchMsg.KEYS: [b"k1"],

1534 FetchMsg.BLOCK_INDEXES: [5],

1535 }

1536 )

@@ -743,12 +1572,12 @@ class TestFinishRequestServerSide:

1572 {

1573 TYPE_KEY: FetchMsg.TYPE,

1574 FetchMsg.KV_REQUEST_ID: "req-1",

746 - FetchMsg.BLOCK_HASHES: [b"k1"],

1575 + FetchMsg.KEYS: [b"k1"],

1576 FetchMsg.BLOCK_INDEXES: [5],

1577 }

1578 )

1579 session.poll()

751 - assert "req-1" in session._server._outbound

1580 + assert _srv_outbound(session, "req-1") is not None

1581

1582 session.finish_request("req-1")

1583

@@ -756,7 +1585,7 @@ class TestFinishRequestServerSide:

1585 assert msg is not None

1586 assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"

1587 assert msg[TransferDoneMsg.SUCCESS] is False

759 - assert "req-1" not in session._server._outbound

1588 + assert _srv_outbound(session, "req-1") is None

1589

1590 def test_with_inflight_defers_then_fires_on_last_transfer(self):

1591 """finish_request with inflight defers; last transfer fires the

@@ -768,7 +1597,7 @@ class TestFinishRequestServerSide:

1597 {

1598 TYPE_KEY: FetchMsg.TYPE,

1599 FetchMsg.KV_REQUEST_ID: "req-1",

771 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

1600 + FetchMsg.KEYS: [b"k1", b"k2"],

1601 FetchMsg.BLOCK_INDEXES: [10, 11],

1602 }

1603 )

@@ -780,8 +1609,8 @@ class TestFinishRequestServerSide:

1609 before = len(conn._sent)

1610 session.finish_request("req-1")

1611 assert len(conn._sent) == before

783 - assert "req-1" in session._server._outbound

784 - assert session._server._outbound["req-1"].finishing

1612 + assert _srv_outbound(session, "req-1") is not None

1613 + assert _srv_outbound(session, "req-1").finishing

1614

1615 # Last inflight settles -> early-fail fires.

1616 tid = next(iter(transport._transfers))

@@ -792,7 +1621,7 @@ class TestFinishRequestServerSide:

1621 assert msg is not None

1622 assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"

1623 assert msg[TransferDoneMsg.SUCCESS] is False

795 - assert "req-1" not in session._server._outbound

1624 + assert _srv_outbound(session, "req-1") is None

1625

1626 def test_full_demand_satisfied_still_sends_success(self):

1627 """finish_request must not override a fully-satisfied transfer:

@@ -803,7 +1632,7 @@ class TestFinishRequestServerSide:

1632 {

1633 TYPE_KEY: FetchMsg.TYPE,

1634 FetchMsg.KV_REQUEST_ID: "req-1",

806 - FetchMsg.BLOCK_HASHES: [b"k1"],

1635 + FetchMsg.KEYS: [b"k1"],

1636 FetchMsg.BLOCK_INDEXES: [10],

1637 }

1638 )

@@ -831,13 +1660,13 @@ class TestFinishRequestServerSide:

1660 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)

1661 # finish_request first — no demand received yet -> defer.

1662 session.finish_request("req-1")

834 - assert "req-1" in session._server._outbound

1663 + assert _srv_outbound(session, "req-1") is not None

1664 # Fetch arrives now: demand fully satisfied by available.

1665 conn.enqueue(

1666 {

1667 TYPE_KEY: FetchMsg.TYPE,

1668 FetchMsg.KV_REQUEST_ID: "req-1",

840 - FetchMsg.BLOCK_HASHES: [b"k1"],

1669 + FetchMsg.KEYS: [b"k1"],

1670 FetchMsg.BLOCK_INDEXES: [10],

1671 }

1672 )

@@ -858,7 +1687,7 @@ class TestFinishRequestServerSide:

1687 {

1688 TYPE_KEY: FetchMsg.TYPE,

1689 FetchMsg.KV_REQUEST_ID: "req-2",

861 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

1690 + FetchMsg.KEYS: [b"k1", b"k2"],

1691 FetchMsg.BLOCK_INDEXES: [10, 11],

1692 }

1693 )

@@ -870,7 +1699,7 @@ class TestFinishRequestServerSide:

1699 session.poll()

1700 msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)

1701 assert msg[TransferDoneMsg.SUCCESS] is False

873 - assert "req-2" not in session._server._outbound

1702 + assert _srv_outbound(session, "req-2") is None

1703

1704 def test_unknown_request_is_noop(self):

1705 session, conn, _ = _make_session()

@@ -891,7 +1720,7 @@ class TestFinishRequestServerSide:

1720 {

1721 TYPE_KEY: FetchMsg.TYPE,

1722 FetchMsg.KV_REQUEST_ID: "req-1",

894 - FetchMsg.BLOCK_HASHES: [b"demand"],

1723 + FetchMsg.KEYS: [b"demand"],

1724 FetchMsg.BLOCK_INDEXES: [5],

1725 }

1726 )

@@ -900,7 +1729,7 @@ class TestFinishRequestServerSide:

1729 # matches demand. Without the shortcut, job 42 sits in _store_jobs

1730 # for _STORE_TIMEOUT_S.

1731 session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42)

903 - assert session._server._outbound["req-1"].pending_job_ids == {42}

1732 + assert _srv_outbound(session, "req-1").pending_job_ids == {42}

1733

1734 session.finish_request("req-1")

1735

@@ -908,7 +1737,7 @@ class TestFinishRequestServerSide:

1737 msg = self._last_transfer_done(conn)

1738 assert msg is not None

1739 assert msg[TransferDoneMsg.SUCCESS] is False

911 - assert "req-1" not in session._server._outbound

1740 + assert _srv_outbound(session, "req-1") is None

1741

1742 # Local store job surfaces on the next poll, success=False.

1743 stores = session.poll().stores

@@ -925,7 +1754,7 @@ class TestFinishRequestServerSide:

1754 {

1755 TYPE_KEY: FetchMsg.TYPE,

1756 FetchMsg.KV_REQUEST_ID: "req-1",

928 - FetchMsg.BLOCK_HASHES: [b"k1"],

1757 + FetchMsg.KEYS: [b"k1"],

1758 FetchMsg.BLOCK_INDEXES: [10],

1759 }

1760 )

@@ -933,7 +1762,7 @@ class TestFinishRequestServerSide:

1762 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7)

1763 # finish_request races with the inflight transfer.

1764 session.finish_request("req-1")

936 - assert "req-1" in session._server._outbound # deferred

1765 + assert _srv_outbound(session, "req-1") is not None # deferred

1766

1767 # Last inflight completes -> _finalize_outbound(success=True) fires.

1768 tid = next(iter(transport._transfers))

@@ -944,7 +1773,7 @@ class TestFinishRequestServerSide:

1773 assert msg is not None

1774 assert msg[TransferDoneMsg.SUCCESS] is True

1775 assert StoreResult(job_id=7, success=True) in stores

947 - assert "req-1" not in session._server._outbound

1776 + assert _srv_outbound(session, "req-1") is None

1777

1778 def test_write_blocks_failure_finalizes_with_failure(self):

1779 """write_blocks returning None must not leave the request hanging.

@@ -962,7 +1791,7 @@ class TestFinishRequestServerSide:

1791 {

1792 TYPE_KEY: FetchMsg.TYPE,

1793 FetchMsg.KV_REQUEST_ID: "req-1",

965 - FetchMsg.BLOCK_HASHES: [b"k1"],

1794 + FetchMsg.KEYS: [b"k1"],

1795 FetchMsg.BLOCK_INDEXES: [10],

1796 }

1797 )

@@ -973,7 +1802,7 @@ class TestFinishRequestServerSide:

1802 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42)

1803

1804 # Outbound was finalized immediately (no other inflight).

976 - assert "req-1" not in session._server._outbound

1805 + assert _srv_outbound(session, "req-1") is None

1806 # Peer notified with success=False.

1807 msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)

1808 assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"

@@ -996,14 +1825,14 @@ class TestFinishRequestServerSide:

1825 {

1826 TYPE_KEY: FetchMsg.TYPE,

1827 FetchMsg.KV_REQUEST_ID: "req-1",

999 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2", b"k3"],

1828 + FetchMsg.KEYS: [b"k1", b"k2", b"k3"],

1829 FetchMsg.BLOCK_INDEXES: [10, 11, 12],

1830 }

1831 )

1832 session.poll()

1833 # Demand registered, no matches yet.

1834 assert session._server._inflight == {}

1006 - outbound = session._server._outbound["req-1"]

1835 + outbound = _srv_outbound(session, "req-1")

1836 assert outbound.remaining == 3

1837 assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"}

1838

@@ -1024,7 +1853,7 @@ class TestFinishRequestServerSide:

1853 assert session._server._inflight == {}

1854 assert outbound.remaining == 2

1855 # Not yet finalized — still 2 blocks demanded.

1027 - assert "req-1" in session._server._outbound

1856 + assert _srv_outbound(session, "req-1") is not None

1857

1858 # Round 2: k2 and k3 arrive together.

1859 session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200)

@@ -1038,7 +1867,7 @@ class TestFinishRequestServerSide:

1867 stores = session.poll().stores

1868 assert StoreResult(job_id=200, success=True) in stores

1869 # _finalize_outbound fired — request gone, peer notified with success.

1041 - assert "req-1" not in session._server._outbound

1870 + assert _srv_outbound(session, "req-1") is None

1871 done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE)

1872 assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1"

1873 assert done[TransferDoneMsg.SUCCESS] is True

@@ -1057,7 +1886,7 @@ class TestFinishRequestServerSide:

1886 {

1887 TYPE_KEY: FetchMsg.TYPE,

1888 FetchMsg.KV_REQUEST_ID: "req-1",

1060 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

1889 + FetchMsg.KEYS: [b"k1", b"k2"],

1890 FetchMsg.BLOCK_INDEXES: [10, 11],

1891 }

1892 )

@@ -1067,7 +1896,7 @@ class TestFinishRequestServerSide:

1896 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100)

1897 assert len(session._server._inflight) == 1

1898 tid_1 = next(iter(session._server._inflight))

1070 - outbound = session._server._outbound["req-1"]

1899 + outbound = _srv_outbound(session, "req-1")

1900 assert outbound.remaining == 2 # decrement happens on completion

1901 assert outbound.finishing is False

1902

@@ -1078,7 +1907,7 @@ class TestFinishRequestServerSide:

1907 assert list(session._server._inflight.keys()) == [tid_1]

1908 # Marked finishing, but NOT finalized yet (transfer_1 still inflight).

1909 assert outbound.finishing is True

1081 - assert "req-1" in session._server._outbound

1910 + assert _srv_outbound(session, "req-1") is not None

1911 done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE]

1912 assert done_msgs == []

1913

@@ -1090,7 +1919,7 @@ class TestFinishRequestServerSide:

1919 stores_first = session.poll().stores

1920 assert StoreResult(job_id=100, success=True) in stores_first

1921 # Outbound state cleaned up; peer notified with success=False.

1093 - assert "req-1" not in session._server._outbound

1922 + assert _srv_outbound(session, "req-1") is None

1923 done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE)

1924 assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1"

1925 assert done[TransferDoneMsg.SUCCESS] is False

@@ -1123,7 +1952,7 @@ class TestBidirectional:

1952 {

1953 TYPE_KEY: FetchMsg.TYPE,

1954 FetchMsg.KV_REQUEST_ID: "req-srv",

1126 - FetchMsg.BLOCK_HASHES: [b"served"],

1955 + FetchMsg.KEYS: [b"served"],

1956 FetchMsg.BLOCK_INDEXES: [7],

1957 }

1958 )

@@ -1178,6 +2007,7 @@ class TestPendingSession:

2007 local_id="local:9000",

2008 transport=transport, # type: ignore[arg-type]

2009 local_block_len=4096,

2010 + local_hash_seed=_DEFAULT_HASH_SEED,

2011 conn=None,

2012 )

2013 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)

@@ -1197,6 +2027,7 @@ class TestPendingSession:

2027 local_id="local:9000",

2028 transport=transport, # type: ignore[arg-type]

2029 local_block_len=4096,

2030 + local_hash_seed=_DEFAULT_HASH_SEED,

2031 conn=None,

2032 )

2033 conn = FakeConnection(peer_id="peer:8000")

@@ -1218,13 +2049,16 @@ class TestPendingSession:

2049 local_id="local:9000",

2050 transport=transport, # type: ignore[arg-type]

2051 local_block_len=4096,

2052 + local_hash_seed=_DEFAULT_HASH_SEED,

2053 conn=None,

2054 )

2055 session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)

2056 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}

2057 + result = session.close()

2058 + assert result.failed_jobs == []

2059 + assert result.failed_req_ids == []

2060 + assert set(result.failed_stores) == {1, 2}

2061 + assert result.failed_serves == []

2062

2063

2064 # ---------------------------------------------------------------------------

@@ -1246,9 +2080,50 @@ class TestDisconnect:

2080 session.request_blocks(1, "req-1", [b"k"], [0])

2081 session.request_blocks(2, "req-2", [b"k"], [0])

2082 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}

2083 + result = session.close()

2084 + assert set(result.failed_jobs) == {1, 2}

2085 + assert set(result.failed_req_ids) == {"req-1", "req-2"}

2086 + assert set(result.failed_stores) == {10}

2087 + assert result.failed_serves == []

2088 +

2089 + def test_send_failure_marks_connection_dead(self):

2090 + """A raising send must mark the connection dead, not silently drop

2091 + the message — otherwise the session lingers alive, is never reaped,

2092 + and in-flight lookups/loads toward the dead peer hang forever."""

2093 + session, conn, _ = _make_session()

2094 + _activate(session, conn)

2095 + assert session.alive

2096 +

2097 + conn.fail_send = True

2098 + # request_blocks flushes a FetchMsg synchronously via _do_send.

2099 + session.request_blocks(1, "req-1", [b"k"], [0])

2100 +

2101 + assert not session.alive

2102 +

2103 + def test_close_surfaces_inflight_lookups(self):

2104 + """close() reports kv_request_ids whose symmetric-P2P probe is still

2105 + unresolved; resolved probes are not reported (their answer is in)."""

2106 + session, conn, _ = _make_session()

2107 + _activate(session, conn)

2108 +

2109 + session.register_lookup("req-hit", b"hA")

2110 + session.register_lookup("req-inflight", b"hB")

2111 + session.flush_pending_lookups()

2112 +

2113 + # Only req-hit is answered; req-inflight stays in flight.

2114 + conn.enqueue(

2115 + {

2116 + TYPE_KEY: LookupRespMsg.TYPE,

2117 + LookupRespMsg.KV_REQUEST_ID: "req-hit",

2118 + LookupRespMsg.KEYS: [b"hA"],

2119 + LookupRespMsg.HITS: [True],

2120 + }

2121 + )

2122 + session.poll()

2123 +

2124 + result = session.close()

2125 + assert result.failed_jobs == []

2126 + assert result.failed_req_ids == ["req-inflight"]

2127

2128

2129 # ---------------------------------------------------------------------------

@@ -1294,7 +2169,7 @@ class TestAdversarial:

2169 {

2170 TYPE_KEY: FetchMsg.TYPE,

2171 FetchMsg.KV_REQUEST_ID: "req-bad",

1297 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

2172 + FetchMsg.KEYS: [b"k1", b"k2"],

2173 FetchMsg.BLOCK_INDEXES: [1],

2174 }

2175 )

@@ -1342,7 +2217,7 @@ class TestDispatchErrorHandling:

2217 {

2218 TYPE_KEY: FetchMsg.TYPE,

2219 FetchMsg.KV_REQUEST_ID: "req-bad",

1345 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

2220 + FetchMsg.KEYS: [b"k1", b"k2"],

2221 FetchMsg.BLOCK_INDEXES: [1],

2222 }

2223 )

@@ -1372,7 +2247,7 @@ class TestDispatchErrorHandling:

2247 {

2248 TYPE_KEY: FetchMsg.TYPE,

2249 FetchMsg.KV_REQUEST_ID: "req-1",

1375 - FetchMsg.BLOCK_HASHES: [b"k1"],

2250 + FetchMsg.KEYS: [b"k1"],

2251 FetchMsg.BLOCK_INDEXES: [0],

2252 }

2253 )

@@ -1395,7 +2270,7 @@ class TestDispatchErrorHandling:

2270 {

2271 TYPE_KEY: FetchMsg.TYPE,

2272 FetchMsg.KV_REQUEST_ID: "req-1",

1398 - FetchMsg.BLOCK_HASHES: [b"k1"],

2273 + FetchMsg.KEYS: [b"k1"],

2274 FetchMsg.BLOCK_INDEXES: [0],

2275 }

2276 )

@@ -1422,7 +2297,7 @@ class TestDispatchErrorHandling:

2297 {

2298 TYPE_KEY: FetchMsg.TYPE,

2299 FetchMsg.KV_REQUEST_ID: "req-1",

1425 - FetchMsg.BLOCK_HASHES: [b"k1"],

2300 + FetchMsg.KEYS: [b"k1"],

2301 FetchMsg.BLOCK_INDEXES: [0],

2302 }

2303 )

@@ -1438,10 +2313,10 @@ class TestDispatchErrorHandling:

2313

2314

2315 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.

2316 + """Per-request `inflight_tids` is the O(1) replacement for the

2317 + previous O(N) scan in `_has_inflight_for`. These tests check that

2318 + every mutation site keeps the set in sync with `_inflight` and

2319 + that the lookup is correct under high fan-out.

2320 """

2321

2322 def test_invariant_holds_through_lifecycle(self):

@@ -1452,17 +2327,14 @@ class TestInflightPerReqInvariant:

2327 _activate(session, conn)

2328

2329 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 - )

2330 + return _srv_total_inflight(session) == len(session._server._inflight)

2331

2332 assert _invariant_holds()

2333

2334 # Two requests, two blocks each, all dispatched in one batch.

2335 session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10)

2336 session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11)

1465 - for kv_id, hashes, indexes in (

2337 + for kv_id, keys, indexes in (

2338 ("req-A", [b"a1", b"a2"], [100, 101]),

2339 ("req-B", [b"b1", b"b2"], [102, 103]),

2340 ):

@@ -1470,7 +2342,7 @@ class TestInflightPerReqInvariant:

2342 {

2343 TYPE_KEY: FetchMsg.TYPE,

2344 FetchMsg.KV_REQUEST_ID: kv_id,

1473 - FetchMsg.BLOCK_HASHES: hashes,

2345 + FetchMsg.KEYS: keys,

2346 FetchMsg.BLOCK_INDEXES: indexes,

2347 }

2348 )

@@ -1493,7 +2365,7 @@ class TestInflightPerReqInvariant:

2365

2366 assert _invariant_holds()

2367 assert not session._server._has_inflight_for("req-A")

1496 - assert "req-A" not in session._server._inflight_per_req # entry was removed

2368 + assert _srv_inflight_count(session, "req-A") == 0 # entry drained

2369 assert session._server._has_inflight_for("req-B")

2370

2371 # Complete req-B; counter must drain to empty.

@@ -1508,7 +2380,7 @@ class TestInflightPerReqInvariant:

2380

2381 assert _invariant_holds()

2382 assert session._server._inflight == {}

1511 - assert session._server._inflight_per_req == {}

2383 + assert _srv_total_inflight(session) == 0

2384

2385 def test_has_inflight_for_correct_with_many_requests(self):

2386 """Populate many inflight xfers across many ids; lookup must

@@ -1524,9 +2396,7 @@ class TestInflightPerReqInvariant:

2396 tid,

2397 _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}),

2398 )

1527 - assert sum(session._server._inflight_per_req.values()) == len(

1528 - session._server._inflight

1529 - )

2399 + assert _srv_total_inflight(session) == len(session._server._inflight)

2400 assert session._server._has_inflight_for("req-0")

2401 assert session._server._has_inflight_for("req-99")

2402 assert not session._server._has_inflight_for("req-missing")

@@ -1539,7 +2409,7 @@ class TestInflightPerReqInvariant:

2409 ]

2410 for tid in tids_50:

2411 session._server._inflight_pop(tid)

1542 - assert "req-50" not in session._server._inflight_per_req

2412 + assert _srv_inflight_count(session, "req-50") == 0

2413 assert not session._server._has_inflight_for("req-50")

2414 # Other ids unaffected.

2415 assert session._server._has_inflight_for("req-49")

@@ -1559,6 +2429,7 @@ class TestConnectMsgValidation:

2429 ConnectMsg.BASE_ADDR: 0x1000,

2430 ConnectMsg.NUM_BLOCKS: 8,

2431 ConnectMsg.BLOCK_LEN: 4096,

2432 + ConnectMsg.HASH_SEED: "0",

2433 }

2434

2435 def test_valid_message_passes(self):

@@ -1600,13 +2471,25 @@ class TestConnectMsgValidation:

2471 with pytest.raises(ValueError, match="block_len"):

2472 ConnectMsg.validate(msg)

2473

2474 + def test_missing_hash_seed(self):

2475 + msg = self._valid_msg()

2476 + del msg[ConnectMsg.HASH_SEED]

2477 + with pytest.raises(ValueError, match="hash_seed"):

2478 + ConnectMsg.validate(msg)

2479 +

2480 + def test_hash_seed_wrong_type(self):

2481 + msg = self._valid_msg()

2482 + msg[ConnectMsg.HASH_SEED] = 12345 # int, not str

2483 + with pytest.raises(ValueError, match="hash_seed"):

2484 + ConnectMsg.validate(msg)

2485 +

2486

2487 class TestFetchMsgValidation:

2488 def _valid_msg(self) -> dict:

2489 return {

2490 TYPE_KEY: FetchMsg.TYPE,

2491 FetchMsg.KV_REQUEST_ID: "req-1",

1609 - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],

2492 + FetchMsg.KEYS: [b"k1", b"k2"],

2493 FetchMsg.BLOCK_INDEXES: [0, 1],

2494 }

2495

vllm/v1/kv_offload/base.py

+15 -2

@@ -6,9 +6,9 @@ Core abstractions for KV cache offloading in vLLM v1.

7 from abc import ABC, abstractmethod

8 from collections.abc import Collection, Iterable, Sequence

9 -from dataclasses import dataclass

9 +from dataclasses import dataclass, field

10 from enum import Enum, auto

11 -from typing import TYPE_CHECKING, Any, NamedTuple, NewType

11 +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeVar

13 import numpy as np

14 import torch

@@ -45,10 +45,23 @@ def get_offload_group_idx(key: OffloadKey) -> int:

45 return int.from_bytes(key[-4:], "big", signed=False)

48 +_T = TypeVar("_T")

49 +

50 +

51 @dataclass

52 class ReqContext:

53 req_id: str

54 kv_transfer_params: dict[str, Any] | None = None

55 + # Per-request scratch space keyed by value type, so a tier can parse

56 + # kv_transfer_params once (in on_new_request) and read the result back

57 + # on later calls for the same request.

58 + _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False)

59 +

60 + def set_state(self, val: Any) -> None:

61 + self._state[type(val)] = val

62 +

63 + def get_state(self, cls: type[_T]) -> _T | None:

64 + return self._state.get(cls)

67 class LookupResult(Enum):

vllm/v1/kv_offload/tiering/p2p/manager.py

+235 -93

@@ -8,6 +8,7 @@ Owns transports and a single bidirectional P2PSession per remote peer.

9 from __future__ import annotations

11 +import os

12 import time

13 import uuid

14 from collections.abc import Iterable, Sequence

@@ -23,12 +24,12 @@ from vllm.v1.kv_offload.base import (

24 OffloadKey,

25 ReqContext,

26 RequestOffloadingContext,

27 + ScheduleEndContext,

28 )

29 from vllm.v1.kv_offload.file_mapper import FileMapper

30 from vllm.v1.kv_offload.tiering.base import (

31 JobMetadata,

32 JobResult,

31 - ScheduleEndContext,

33 SecondaryTierManager,

34 )

35 from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport

@@ -37,6 +38,7 @@ from vllm.v1.kv_offload.tiering.p2p.session import P2PSession

39 if TYPE_CHECKING:

40 from vllm.v1.kv_offload.base import OffloadingSpec

41 + from vllm.v1.kv_offload.tiering.base import ParentManager

42 from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection

44 logger = init_logger(__name__)

@@ -59,24 +61,112 @@ _SHUTDOWN_DRAIN_TIMEOUT_S = 3.0

61 _DRAIN_SLEEP_S = 0.001

62 -def _prefill_params(kv_params: dict | None) -> dict | None:

63 - """Return the ``prefill`` sub-dict, or None if absent.

64 +def _remote_prefiller_params(kv_params: dict | None) -> dict | None:

65 + """Return the ``remote_prefiller`` sub-dict, or None if absent.

66 +

67 + Set on decoder requests to name the remote prefiller they pull from;

68 + carries kv_request_id, remote_host, remote_port.

69 + """

70 + if not kv_params:

71 + return None

72 + return kv_params.get("remote_prefiller")

73 +

74 +

75 +def _remote_decoder_params(kv_params: dict | None) -> dict | None:

76 + """Return the ``remote_decoder`` sub-dict, or None if absent.

65 - Set on decoder requests; carries kv_request_id, remote_host, remote_port.

78 + Set on prefiller requests to name the remote decoder they serve;

79 + carries kv_request_id.

80 """

81 if not kv_params:

82 return None

69 - return kv_params.get("prefill")

83 + return kv_params.get("remote_decoder")

72 -def _decode_params(kv_params: dict | None) -> dict | None:

73 - """Return the ``decode`` sub-dict, or None if absent.

86 +def _remote_kv_source_params(kv_params: dict | None) -> dict | None:

87 + """Return the ``remote_kv_source`` sub-dict, or None if absent.

75 - Set on prefiller requests; carries kv_request_id.

89 + Set on symmetric-P2P consumer requests to name the remote source they

90 + pull from; carries kv_request_id, remote_host, remote_port.

91 """

92 if not kv_params:

93 return None

79 - return kv_params.get("decode")

94 + return kv_params.get("remote_kv_source")

95 +

96 +

97 +def _peer_id_from_params(role_params: dict) -> str | None:

98 + """Build ``host:port`` peer_id from a role-scoped sub-dict, or None."""

99 + host = role_params.get("remote_host")

100 + port = role_params.get("remote_port")

101 + if host and port:

102 + return f"{host}:{port}"

103 + return None

104 +

105 +

106 +@dataclass(slots=True)

107 +class P2PSourceInfo:

108 + """Consumer side: this request fetches from a remote (prefiller or peer)."""

109 +

110 + kv_request_id: str

111 + peer_id: str

112 + do_probe: bool # False for remote_prefiller (PD), True for remote_kv_source

113 +

114 +

115 +@dataclass(slots=True)

116 +class P2PDestInfo:

117 + """Producer side: a remote fetches this request's blocks from us.

118 +

119 + ``kv_request_id`` is None when the ``remote_decoder`` block is present

120 + but malformed (no id); the block's presence still marks the request as

121 + remote-decode, so submit_store must fail rather than store locally.

122 + """

123 +

124 + kv_request_id: str | None

125 +

126 +

127 +def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None:

128 + """Parse the consumer sub-dict (PD ``remote_prefiller`` or symmetric

129 + ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete."""

130 + role = _remote_prefiller_params(kv_params)

131 + do_probe = False

132 + if role is None:

133 + role = _remote_kv_source_params(kv_params)

134 + do_probe = True

135 + if not role:

136 + return None

137 + peer_id = _peer_id_from_params(role)

138 + kv_request_id = role.get("kv_request_id")

139 + if peer_id is None or not kv_request_id:

140 + return None

141 + return P2PSourceInfo(

142 + kv_request_id=kv_request_id,

143 + peer_id=peer_id,

144 + do_probe=do_probe,

145 + )

146 +

147 +

148 +def _parse_dest(kv_params: dict | None) -> P2PDestInfo | None:

149 + """Parse the producer ``remote_decoder`` sub-dict into a ``P2PDestInfo``,

150 + or None if the block is absent (not a remote-decode request)."""

151 + role = _remote_decoder_params(kv_params)

152 + if role is None:

153 + return None

154 + return P2PDestInfo(kv_request_id=role.get("kv_request_id") or None)

155 +

156 +

157 +def _annotate_req_context(req_context: ReqContext) -> None:

158 + """Parse kv_transfer_params once and cache the P2P routing state.

159 +

160 + Called from ``on_new_request``; later calls for the same request read

161 + the cached ``P2PSourceInfo``/``P2PDestInfo`` via ``get_state`` instead

162 + of re-parsing.

163 + """

164 + source = _parse_source(req_context.kv_transfer_params)

165 + if source is not None:

166 + req_context.set_state(source)

167 + dest = _parse_dest(req_context.kv_transfer_params)

168 + if dest is not None:

169 + req_context.set_state(dest)

170

171

172 @dataclass

@@ -157,6 +247,20 @@ class P2PSecondaryTierManager(SecondaryTierManager):

247 **kwargs: Reserved for future tier-specific options.

248 """

249 super().__init__(offloading_spec, primary_kv_view, tier_type)

250 + # Block hashes chain from NONE_HASH, seeded from PYTHONHASHSEED

251 + # (see init_none_hash in v1/core/kv_cache_utils.py). Peers with

252 + # different seeds compute different hashes for identical content, so

253 + # lookups silently miss and no KV crosses the wire. Require it here so

254 + # a misconfigured P2P instance fails at startup rather than degrading

255 + # silently; the value is also verified against each peer on handshake.

256 + hash_seed = os.getenv("PYTHONHASHSEED")

257 + if hash_seed is None:

258 + raise ValueError(

259 + "PYTHONHASHSEED must be set for P2P KV offload so that block "

260 + "hashes match across instances. Set it to a fixed value (e.g. "

261 + "PYTHONHASHSEED=0) on every P2P peer."

262 + )

263 + self._hash_seed = hash_seed

264 if host is None:

265 host = envs.VLLM_P2P_SIDE_CHANNEL_HOST

266 if port is None:

@@ -210,6 +314,12 @@ class P2PSecondaryTierManager(SecondaryTierManager):

314 # kv_request_ids that hit a transport/session failure; On load lookup()

315 # rejects them so the request falls back to local prefill.

316 self._failed_req_ids: set[str] = set()

317 + # Synthetic lookup ctxs from reaped sessions still owing a

318 + # ``parent.on_request_finished`` (the session's failed_serves). The

319 + # dead session had no parent handle at teardown; these are flushed

320 + # at the top of the next ``serve_external_requests`` where the

321 + # handle is valid.

322 + self._failed_serve_ctxs: list[ReqContext] = []

323

324 # ------------------------------------------------------------------

325 # SecondaryTierManager interface

@@ -217,66 +327,81 @@ class P2PSecondaryTierManager(SecondaryTierManager):

327

328 @override

329 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

220 - prefill = _prefill_params(req_context.kv_transfer_params)

221 - if (

222 - not prefill

223 - or not prefill.get("remote_host")

224 - or not prefill.get("remote_port")

225 - or not prefill.get("kv_request_id")

226 - ):

330 + source = req_context.get_state(P2PSourceInfo)

331 + if source is None:

332 return LookupResult.MISS

228 -

229 - kv_request_id = prefill["kv_request_id"]

230 - if kv_request_id in self._failed_req_ids:

333 + if source.kv_request_id in self._failed_req_ids:

334 return LookupResult.MISS

335 +

336 + # Symmetric-P2P consumer (``remote_kv_source`` sub-dict): probe the

337 + # peer asynchronously. First call registers the (kv_request_id,

338 + # key) entry and returns RETRY; flush_pending_lookups()

339 + # in on_schedule_end batches the LookupMsg; a later step's

340 + # lookup() returns HIT/MISS once LookupRespMsg has arrived.

341 + # PD path (``remote_prefiller`` sub-dict only) keeps the eager HIT.

342 + if source.do_probe:

343 + session = self._sessions.get(source.peer_id)

344 + if session is None:

345 + return LookupResult.MISS

346 + result = session.register_lookup(source.kv_request_id, key)

347 + if result is True:

348 + return LookupResult.HIT

349 + if result is False:

350 + return LookupResult.MISS

351 + return LookupResult.RETRY

352 +

353 + # PD consumer (we are the decoder): all kv blocks should be on the

354 + # prefiller side. Return HIT immediately.

355 return LookupResult.HIT

356

357 @override

358 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

236 - """Open the outbound session toward the producer if needed.

237 -

238 - On the decoder side (``prefill`` set), open a session toward the

239 - producer at remote_host:remote_port so submit_load can issue

240 - FetchMsg as soon as it fires. On the prefiller side, sessions

241 - are created when the consumer's inbound connection arrives in

242 - _accept_new_peers — submit_store no longer pre-creates anything.

359 + """Parse kv_transfer_params once and open the outbound session.

360 +

361 + Parses the P2P routing state onto ``req_context`` (cached for the

362 + later lookup/submit/finish calls). On the consumer side

363 + (``remote_prefiller`` for PD or ``remote_kv_source`` for symmetric

364 + P2P), open a session toward the producer at remote_host:remote_port

365 + so submit_load can issue FetchMsg as soon as it fires. On the

366 + prefiller side, sessions are created when the consumer's inbound

367 + connection arrives in _accept_new_peers — submit_store no longer

368 + pre-creates anything.

369 """

244 - prefill = _prefill_params(req_context.kv_transfer_params)

245 - if prefill:

246 - peer_id = self._remote_id_from_params(prefill)

247 - if peer_id:

248 - self._get_or_create_session(peer_id)

370 + _annotate_req_context(req_context)

371 + source = req_context.get_state(P2PSourceInfo)

372 + if source is not None:

373 + self._get_or_create_session(source.peer_id)

374 return RequestOffloadingContext()

375

376 @override

377 def on_request_finished(self, req_context: ReqContext) -> None:

378 """Cancels pending loads and prunes session-scoped state.

379

255 - Decoder side (``prefill`` set): looks up the session by peer_id

256 - because the producer's address is what addresses the client-role

257 - load to cancel. Prefiller side (``decode`` set): looks up via

258 - kv_request_id because peer_id is no longer carried on store-time

380 + Consumer side (``remote_prefiller`` for PD or ``remote_kv_source``

381 + for symmetric-P2P): looks up the session by peer_id because the

382 + producer's address is what addresses the client-role load to

383 + cancel; also drops any pending symmetric-P2P lookup state via

384 + ``session.finish_request``.

385 + Prefiller side (``remote_decoder`` set): looks up via kv_request_id

386 + because peer_id is no longer carried on store-time

387 kv_transfer_params; if a session has bound the id, finish it. If

388 no session has bound the id yet, this is a no-op: parked batches

389 in `_unbound_stores` are left in place and cleaned up only by

390 `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.

391 """

264 - kv_params = req_context.kv_transfer_params

265 - if not kv_params:

266 - return

267 - prefill = _prefill_params(kv_params)

268 - decode = _decode_params(kv_params)

269 - kv_request_id = (prefill or decode or {}).get("kv_request_id")

392 + source = req_context.get_state(P2PSourceInfo)

393 + dest = req_context.get_state(P2PDestInfo)

394 + kv_request_id = source.kv_request_id if source is not None else None

395 + if kv_request_id is None and dest is not None:

396 + kv_request_id = dest.kv_request_id

397 if not kv_request_id:

398 return

399 self._failed_req_ids.discard(kv_request_id)

400

274 - if prefill:

275 - peer_id = self._remote_id_from_params(prefill)

276 - if peer_id:

277 - session = self._sessions.get(peer_id)

278 - if session is not None:

279 - session.finish_request(kv_request_id)

401 + if source is not None:

402 + session = self._sessions.get(source.peer_id)

403 + if session is not None:

404 + session.finish_request(kv_request_id)

405 return

406

407 # Prefiller-side finish: identify the session via kv_request_id.

@@ -293,24 +418,24 @@ class P2PSecondaryTierManager(SecondaryTierManager):

418

419 assert len(keys) == len(block_ids)

420

296 - kv_params = job_metadata.req_context.kv_transfer_params

297 - decode = _decode_params(kv_params)

421 + dest = job_metadata.req_context.get_state(P2PDestInfo)

422 logger.debug(

299 - "P2P %s: submit_store ENTRY job_id=%d blocks=%d decode=%s kv_request_id=%s",

423 + "P2P %s: submit_store ENTRY job_id=%d blocks=%d "

424 + "remote_decoder=%s kv_request_id=%s",

425 self._local_id,

426 job_id,

427 len(block_ids),

303 - decode is not None,

304 - (decode or {}).get("kv_request_id"),

428 + dest is not None,

429 + dest.kv_request_id if dest is not None else None,

430 )

306 - # Absent ``decode`` block => not a remote-decode request: succeed

307 - # locally without parking. An empty/malformed dict is still a

308 - # remote-decode signal and must fail the missing-id check below.

309 - if decode is None:

431 + # Absent ``remote_decoder`` block => not a remote-decode request:

432 + # succeed locally without parking. An empty/malformed dict is still

433 + # a remote-decode signal and must fail the missing-id check below.

434 + if dest is None:

435 self._finished_jobs.append(JobResult(job_id=job_id, success=True))

436 return

437

313 - kv_request_id = decode.get("kv_request_id")

438 + kv_request_id = dest.kv_request_id

439 if not kv_request_id:

440 logger.warning(

441 "P2P %s: submit_store missing kv_request_id",

@@ -350,32 +475,26 @@ class P2PSecondaryTierManager(SecondaryTierManager):

475 keys = list(job_metadata.keys)

476 block_ids = job_metadata.block_ids

477

353 - prefill = _prefill_params(job_metadata.req_context.kv_transfer_params)

478 + source = job_metadata.req_context.get_state(P2PSourceInfo)

479 logger.debug(

480 "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s",

481 self._local_id,

482 job_id,

483 len(block_ids),

359 - (prefill or {}).get("kv_request_id"),

360 - self._remote_id_from_params(prefill or {}),

484 + source.kv_request_id if source is not None else None,

485 + source.peer_id if source is not None else None,

486 )

362 - if (

363 - not prefill

364 - or not prefill.get("remote_host")

365 - or not prefill.get("remote_port")

366 - or not prefill.get("kv_request_id")

367 - ):

487 + if source is None:

488 logger.debug(

369 - "P2P %s: submit_load job_id=%d FAILED missing prefill params",

489 + "P2P %s: submit_load job_id=%d FAILED missing consumer params",

490 self._local_id,

491 job_id,

492 )

493 self._finished_jobs.append(JobResult(job_id=job_id, success=False))

494 return

495

376 - kv_request_id = prefill["kv_request_id"]

377 - peer_id = self._remote_id_from_params(prefill)

378 - assert peer_id is not None # guaranteed by prefill checks above

496 + kv_request_id = source.kv_request_id

497 + peer_id = source.peer_id

498

499 if not keys:

500 logger.debug(

@@ -442,10 +561,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

561 warned = False

562 while True:

563 self._poll_once()

445 - pending = any(

446 - s._client._inbound or s._server._inflight

447 - for s in self._sessions.values()

448 - )

564 + pending = any(s.has_pending_work for s in self._sessions.values())

565 if not pending:

566 return

567 if not warned and time.monotonic() - start > 5.0:

@@ -456,31 +572,45 @@ class P2PSecondaryTierManager(SecondaryTierManager):

572 warned = True

573 time.sleep(_DRAIN_SLEEP_S)

574

575 + @override

576 + def serve_external_requests(self, parent: ParentManager) -> None:

577 + """Serve inbound peer lookups against the tiering manager.

578 +

579 + Called once per scheduler step (before this tier's

580 + ``on_schedule_end``) with a ``parent`` handle valid only for the

581 + duration of the call — the sole window in which the P2P server

582 + role may query the tiering manager. First release bookkeeping for

583 + the failed serves left by a reaped session, then let every live

584 + session resolve its enqueued inbound LookupMsgs.

585 + """

586 + if self._failed_serve_ctxs:

587 + for ctx in self._failed_serve_ctxs:

588 + parent.on_request_finished(ctx)

589 + self._failed_serve_ctxs = []

590 + for session in self._sessions.values():

591 + session.serve_external_requests(parent)

592 +

593 @override

594 def on_schedule_end(self, context: ScheduleEndContext) -> None:

461 - return

595 + # Flush any p2p lookups aggregated during this step.

596 + # One LookupMsg per (peer, kv_request_id) with unsent entries;

597 + # send-gating happens inside the session if not yet ready.

598 + for session in self._sessions.values():

599 + session.flush_pending_lookups()

600

601 # ------------------------------------------------------------------

602 # Internal

603 # ------------------------------------------------------------------

604

467 - @staticmethod

468 - def _remote_id_from_params(role_params: dict) -> str | None:

469 - """Build peer_id from a role-scoped sub-dict (``prefill``/``p2p``)."""

470 - host = role_params.get("remote_host")

471 - port = role_params.get("remote_port")

472 - if host and port:

473 - return f"{host}:{port}"

474 - return None

475 -

605 def _get_or_create_session(self, peer_id: str) -> P2PSession:

606 """Return the existing session for peer_id, or open one outbound.

607

479 - Decoder-side helper for on_new_request: when ``prefill`` is set,

480 - the consumer must reach the producer at peer_id. If we already

481 - have a session toward that peer (from a prior load or a

482 - peer-initiated inbound), reuse it; otherwise open an outbound

483 - ControlConnection and build a connected session.

608 + Consumer-side helper for on_new_request: when ``remote_prefiller``

609 + (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer must reach the

610 + producer at peer_id. If we already have a session toward that

611 + peer (from a prior load or a peer-initiated inbound), reuse it;

612 + otherwise open an outbound ControlConnection and build a

613 + connected session.

614 """

615 session = self._sessions.get(peer_id)

616 if session is not None:

@@ -491,6 +621,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

621 local_id=self._local_id,

622 transport=self._data,

623 local_block_len=self._data.block_len,

624 + local_hash_seed=self._hash_seed,

625 conn=conn,

626 )

627 self._sessions[peer_id] = session

@@ -512,6 +643,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

643 local_id=self._local_id,

644 transport=self._data,

645 local_block_len=self._data.block_len,

646 + local_hash_seed=self._hash_seed,

647 conn=conn,

648 )

649 logger.info(

@@ -546,12 +678,19 @@ class P2PSecondaryTierManager(SecondaryTierManager):

678 ]

679 for kid in stale_kv_ids:

680 del self._kv_to_session[kid]

549 - failed_loads, failed_stores = session.close()

550 - for job_id, kv_request_id in failed_loads:

681 + close_result = session.close()

682 + for job_id in close_result.failed_jobs:

683 self._finished_jobs.append(JobResult(job_id=job_id, success=False))

552 - self._failed_req_ids.add(kv_request_id)

553 - for job_id in failed_stores:

684 + for job_id in close_result.failed_stores:

685 self._finished_jobs.append(JobResult(job_id=job_id, success=False))

686 + # Fail every client-side request (in-flight loads plus unresolved

687 + # symmetric-P2P probes) toward the dead peer so lookup() returns

688 + # MISS (local prefill) instead of RETRY forever — even if a fresh

689 + # session to the same peer is later opened by another request.

690 + self._failed_req_ids.update(close_result.failed_req_ids)

691 + # Release the TieringManager's per-request bookkeeping for the

692 + # dead session's synthetic lookups on the next serve_external_requests.

693 + self._failed_serve_ctxs.extend(close_result.failed_serves)

694 self._data.remove_remote_peer(pid)

695 logger.warning("P2P %s: peer %s down", self._local_id, pid)

696

@@ -648,6 +787,9 @@ class P2PSecondaryTierManager(SecondaryTierManager):

787 def shutdown(self) -> None:

788 self._drain_inflight_for_shutdown()

789 for session in self._sessions.values():

790 + # Orphan ctxs from close() are intentionally dropped: the manager

791 + # is being torn down, so there is no next serve_external_requests

792 + # to flush them and no TieringManager left to release.

793 session.close()

794 self._sessions.clear()

795 self._kv_to_session.clear()

vllm/v1/kv_offload/tiering/p2p/session/__init__.py

@@ -3,6 +3,7 @@

3 from vllm.v1.kv_offload.tiering.p2p.session.session import (

4 LoadResult,

5 P2PSession,

6 + SessionCloseResult,

7 SessionPollResult,

8 StoreResult,

9 )

@@ -10,6 +11,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import (

11 __all__ = [

12 "LoadResult",

13 "P2PSession",

14 + "SessionCloseResult",

15 "SessionPollResult",

16 "StoreResult",

17 ]

vllm/v1/kv_offload/tiering/p2p/session/client.py

+342 -36

@@ -11,16 +11,19 @@ callback injected by the coordinator (which gates on ConnectAck).

12 from __future__ import annotations

14 +import enum

15 import time

16 from collections.abc import Callable, Sequence

16 -from dataclasses import dataclass

17 +from dataclasses import dataclass, field

18 from typing import TYPE_CHECKING, NamedTuple

20 from vllm.logger import init_logger

21 +from vllm.v1.kv_offload.base import OffloadKey

22 from vllm.v1.kv_offload.tiering.p2p.session.protocol import (

23 TYPE_KEY,

24 AbortFetchMsg,

25 FetchMsg,

26 + LookupMsg,

27 )

29 if TYPE_CHECKING:

@@ -32,16 +35,60 @@ _LOAD_TIMEOUT_S = 30.0

35 _ABORT_ACK_TIMEOUT_S = 10.0

38 +class ClientPhase(enum.Enum):

39 + """Lifecycle of a request's client-side lookup/fetch signalling.

40 +

41 + Advances monotonically. Only ``finish`` reads it, to decide

42 + whether a terminal empty FetchMsg is owed to release the peer's

43 + lookup state (owed only from ``PROBING``: a LookupMsg went out but no

44 + FetchMsg has since closed the peer's lookup phase).

45 + """

46 +

47 + REGISTERED = enum.auto() # keys in probes/unsent, nothing sent yet

48 + PROBING = enum.auto() # LookupMsg flushed, awaiting responses

49 + FETCH_SENT = enum.auto() # FetchMsg sent (real or terminal empty)

50 +

51 +

52 @dataclass

36 -class _InboundRequestState:

37 - """Client-role state for a single load request."""

53 +class _InboundLoadState:

54 + """Client-role state for a single in-flight load request.

55 +

56 + Lives on ``_ClientRequestState.load`` for the duration of a fetch;

57 + the owning kv_request_id is the dict key, so it isn't stored here.

58 + """

60 job_id: int # opaque ID assigned by the manager to this load request

40 - kv_request_id: str

61 submitted_at: float

62 aborted_at: float | None = None

65 +@dataclass

66 +class _ClientRequestState:

67 + """Per-kv_request_id client-side state.

68 +

69 + One entry per kv_request_id we're driving. Lookup-phase fields are

70 + used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave

71 + ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An

72 + entry is dropped once every field is idle — see ``ClientRole._maybe_prune``.

73 + """

74 +

75 + # -- Lookup phase (symmetric P2P only; untouched for PD) --

76 + # Probe outcome per OffloadKey: None while in-flight (registered/sent

77 + # but unresolved), True/False once a LookupRespMsg lands. There is no

78 + # timeout — finish (via finish_request) is guaranteed after

79 + # the request's lookup() calls and clears every probe, so an

80 + # unanswered probe simply stays None until then.

81 + probes: dict[OffloadKey, bool | None] = field(default_factory=dict)

82 + # OffloadKeys registered but not yet flushed onto the wire. Drained and

83 + # cleared by the next flush_pending_lookups.

84 + unsent: list[OffloadKey] = field(default_factory=list)

85 +

86 + # Monotonic lookup/fetch signalling phase; see ``ClientPhase``.

87 + phase: ClientPhase = ClientPhase.REGISTERED

88 + # Set while a fetch is in flight; cleared on completion/abort/timeout.

89 + load: _InboundLoadState | None = None

90 +

91 +

92 class LoadResult(NamedTuple):

93 """Result from a session poll, client side."""

@@ -50,6 +97,21 @@ class LoadResult(NamedTuple):

97 success: bool

100 +class ClientCloseResult(NamedTuple):

101 + """What the manager must fail when a peer session is torn down.

102 +

103 + ``failed_jobs`` is the ``job_id`` of every load still in flight; each

104 + becomes a ``JobResult(success=False)``. ``failed_req_ids`` is every

105 + kv_request_id whose lookup() would otherwise defer forever on the dead

106 + peer — the in-flight loads plus any request holding an unresolved

107 + symmetric-P2P probe. ``failed_jobs`` is the subset of ``failed_req_ids``

108 + that had a load job.

109 + """

110 +

111 + failed_jobs: list[int]

112 + failed_req_ids: list[str]

113 +

114 +

115 class ClientRole:

116 """Client-side load state machine for one peer session.

117

@@ -61,9 +123,55 @@ class ClientRole:

123 def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None:

124 self._peer_id = peer_id

125 self._send = send

64 - self._inbound: dict[str, _InboundRequestState] = {}

126 + # All per-kv_request_id state lives here. Entries are created

127 + # lazily by request_blocks / register_lookup and dropped by

128 + # _maybe_prune once every field is idle.

129 + self._requests: dict[str, _ClientRequestState] = {}

130 + # kv_request_ids with unsent lookup keys for the next flush to

131 + # visit — the work-list that keeps flush_pending_lookups from

132 + # scanning every request each scheduler step. Mirrors the server's

133 + # _serve_pending. Populated by register_lookup, drained by

134 + # flush_pending_lookups, and discarded on finish/close.

135 + self._flush_pending: set[str] = set()

136 + # kv_request_ids with a fetch in flight (``st.load is not None``) —

137 + # the work-list collect_results walks for timeouts, and the

138 + # has_active_loads predicate, instead of scanning every request.

139 + # Kept in exact sync with ``st.load``: armed in request_blocks,

140 + # discarded wherever load is cleared, and cleared on close.

141 + self._active_loads: set[str] = set()

142 self._completed_loads: list[LoadResult] = []

143

144 + # ------------------------------------------------------------------

145 + # State helpers

146 + # ------------------------------------------------------------------

147 +

148 + def _get_or_create_request(self, kv_request_id: str) -> _ClientRequestState:

149 + """Get or create the state entry for a kv_request_id."""

150 + st = self._requests.get(kv_request_id)

151 + if st is None:

152 + st = _ClientRequestState()

153 + self._requests[kv_request_id] = st

154 + return st

155 +

156 + def _maybe_prune(self, kv_request_id: str) -> None:

157 + """Drop the entry once it holds no live load or lookup state.

158 +

159 + The sticky ``phase`` is only read by ``finish``. A probe

160 + clears when its fetch is issued (``request_blocks``) or when the

161 + request finishes (``finish``/``close``); in the former case

162 + ``load`` is set and keeps the entry alive, in the latter the phase

163 + is no longer needed — so dropping on emptiness never loses a phase

164 + still in use.

165 + """

166 + st = self._requests.get(kv_request_id)

167 + if st is not None and st.load is None and not st.probes and not st.unsent:

168 + del self._requests[kv_request_id]

169 +

170 + @property

171 + def has_active_loads(self) -> bool:

172 + """True if any kv_request_id has a fetch in flight."""

173 + return bool(self._active_loads)

174 +

175 # ------------------------------------------------------------------

176 # Public API

177 # ------------------------------------------------------------------

@@ -72,7 +180,7 @@ class ClientRole:

180 self,

181 job_id: JobId,

182 kv_request_id: str,

75 - keys: Sequence[bytes],

183 + keys: Sequence[OffloadKey],

184 block_ids: Sequence[int],

185 send_ready: bool,

186 ) -> None:

@@ -86,44 +194,101 @@ class ClientRole:

194 len(block_ids),

195 send_ready,

196 )

89 - self._inbound[kv_request_id] = _InboundRequestState(

197 + st = self._get_or_create_request(kv_request_id)

198 + st.load = _InboundLoadState(

199 job_id=job_id,

91 - kv_request_id=kv_request_id,

200 submitted_at=time.monotonic(),

201 )

202 + self._active_loads.add(kv_request_id)

203 + st.phase = ClientPhase.FETCH_SENT

204 self._send(

205 {

206 TYPE_KEY: FetchMsg.TYPE,

207 FetchMsg.KV_REQUEST_ID: kv_request_id,

98 - FetchMsg.BLOCK_HASHES: list(keys),

208 + FetchMsg.KEYS: list(keys),

209 FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids],

210 }

211 )

212 + # Issuing the fetch ends this request's lookup phase, so drop all

213 + # probe state. Once the peer serves this fetch both sides unpin, so

214 + # the producer may evict the block; a stale cached True would

215 + # otherwise let a re-scheduled lookup() return HIT without

216 + # re-probing, pointing at a block the producer no longer holds.

217 + # Clearing forces a fresh LookupMsg on re-schedule so the producer

218 + # answers from current state. For a symmetric-P2P request (probes

219 + # populated) every fetched block was a confirmed HIT; a PD-only load

220 + # never probes, so probes is empty and the clear is a no-op.

221 + if st.probes:

222 + assert all(st.probes.get(key) is True for key in keys)

223 + st.probes.clear()

224 +

225 + def finish(self, kv_request_id: str) -> None:

226 + """Finish a request: abort any in-flight load and release lookup state.

227 +

228 + Called from the session's ``finish_request``. The two branches are

229 + mutually exclusive: ``load`` is set only by ``request_blocks``, which

230 + also advances ``phase`` to ``FETCH_SENT``, and nothing moves it back

231 + to ``PROBING`` — so a fetch in flight never coexists with the

232 + ``PROBING`` phase.

233

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:

234 + - Fetch in flight (``load`` set, phase ``FETCH_SENT``): send an

235 + AbortFetchMsg unless the load is already aborting, then drop it.

236 + - Outstanding lookups (``PROBING``): a LookupMsg was flushed but no

237 + FetchMsg has closed the peer's lookup phase. Every FetchMsg the

238 + server receives in p2p mode is its "request finished" signal (it

239 + releases lookup state and fires ``cb.finish_request``); when the

240 + client's lookups all missed no FetchMsg is otherwise sent, so emit

241 + a terminal empty one purely to trigger those semantics. In

242 + ``REGISTERED`` the peer never received a LookupMsg and in

243 + ``FETCH_SENT`` a FetchMsg already closed the phase, so neither owes

244 + a terminal FetchMsg.

245 +

246 + Then drop all probe/lookup state and prune the entry.

247 + """

248 + st = self._requests.get(kv_request_id)

249 + if st is None:

250 + return

251 + if st.load is not None:

252 + if st.load.aborted_at is None:

253 + self._send(

254 + {

255 + TYPE_KEY: AbortFetchMsg.TYPE,

256 + AbortFetchMsg.KV_REQUEST_ID: kv_request_id,

257 + }

258 + )

259 + st.load = None

260 + self._active_loads.discard(kv_request_id)

261 + elif st.phase is ClientPhase.PROBING:

262 + st.phase = ClientPhase.FETCH_SENT

263 self._send(

264 {

109 - TYPE_KEY: AbortFetchMsg.TYPE,

110 - AbortFetchMsg.KV_REQUEST_ID: kv_request_id,

265 + TYPE_KEY: FetchMsg.TYPE,

266 + FetchMsg.KV_REQUEST_ID: kv_request_id,

267 + FetchMsg.KEYS: [],

268 + FetchMsg.BLOCK_INDEXES: [],

269 }

270 )

271 + st.probes.clear()

272 + st.unsent.clear()

273 + self._flush_pending.discard(kv_request_id)

274 + self._maybe_prune(kv_request_id)

275

276 def on_transfer_done(self, kv_request_id: str, success: bool) -> None:

277 """Handle a TransferDoneMsg from the peer."""

116 - req = self._inbound.pop(kv_request_id, None)

117 - if req is not None:

278 + st = self._requests.get(kv_request_id)

279 + if st is not None and st.load is not None:

280 self._completed_loads.append(

281 LoadResult(

120 - job_id=req.job_id,

282 + job_id=st.load.job_id,

283 kv_request_id=kv_request_id,

284 success=success,

285 )

286 )

287 + st.load = None

288 + self._active_loads.discard(kv_request_id)

289 + self._maybe_prune(kv_request_id)

290 else:

126 - # No matching _inbound entry: either a duplicate

291 + # No matching in-flight load: either a duplicate

292 # transfer_done from the peer (protocol violation) or a

293 # benign race with a local cancel/abort/timeout that

294 # already popped the entry. We don't track terminated ids,

@@ -137,23 +302,26 @@ class ClientRole:

302

303 def on_abort_ack(self, kv_request_id: str) -> None:

304 """Handle an AbortAckMsg from the peer."""

140 - req = self._inbound.pop(kv_request_id, None)

141 - if req is not None:

305 + st = self._requests.get(kv_request_id)

306 + if st is not None and st.load is not None:

307 logger.warning(

308 "P2PSession %s: load request %s (job_id=%d) timed out; "

309 "load job completed with failure. If this recurs, ensure "

310 "PYTHONHASHSEED is set to the same value on all nodes.",

311 self._peer_id,

312 kv_request_id,

148 - req.job_id,

313 + st.load.job_id,

314 )

315 self._completed_loads.append(

316 LoadResult(

152 - job_id=req.job_id,

317 + job_id=st.load.job_id,

318 kv_request_id=kv_request_id,

319 success=False,

320 )

321 )

322 + st.load = None

323 + self._active_loads.discard(kv_request_id)

324 + self._maybe_prune(kv_request_id)

325 else:

326 # See on_transfer_done: same ambiguity (duplicate ack

327 # vs. raced with local cancel/timeout that already popped).

@@ -164,19 +332,141 @@ class ClientRole:

332 kv_request_id,

333 )

334

335 + # ------------------------------------------------------------------

336 + # Symmetric-P2P lookup (do_p2p_fetch=true)

337 + # ------------------------------------------------------------------

338 +

339 + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None:

340 + """Register or resolve one (kv_request_id, key) probe.

341 +

342 + Idempotent across scheduler steps:

343 + - First call: creates a pending entry, returns None.

344 + - Subsequent calls while in-flight: returns None.

345 + - Once a LookupRespMsg has resolved the entry: returns the cached

346 + bool result on every call without popping it.

347 +

348 + A resolved entry is retained until its fetch is issued

349 + (``request_blocks`` pops it) or the request finishes

350 + (``finish`` clears all entries for the id). A request's

351 + block set can be re-probed across steps, so popping on read would

352 + make a repeat probe of an already-resolved key look brand-new and

353 + re-queue it, emitting a redundant LookupMsg for an answer we

354 + already hold. Keeping the entry until fetch makes repeat probes

355 + free; clearing it at fetch forces a fresh probe if the request is

356 + re-scheduled, since the block is unpinned once served.

357 + """

358 + st = self._get_or_create_request(kv_request_id)

359 + okey = OffloadKey(key)

360 + if okey in st.probes:

361 + return st.probes[okey]

362 + st.probes[okey] = None

363 + st.unsent.append(okey)

364 + self._flush_pending.add(kv_request_id)

365 + logger.debug(

366 + "P2P LOOKUP client %s: REGISTER kv_request_id=%s key=%s (unsent=%d)",

367 + self._peer_id,

368 + kv_request_id,

369 + key.hex()[:16],

370 + len(st.unsent),

371 + )

372 + return None

373 +

374 + def flush_pending_lookups(self) -> None:

375 + """Send a LookupMsg for each kv_request_id with unsent entries.

376 +

377 + Called once per scheduler step from the manager's

378 + ``on_schedule_end()``. A request's block set may be discovered

379 + across several scheduler steps, so more than one LookupMsg can

380 + go out per kv_request_id — one per step that registered new

381 + keys. register_lookup() de-dups in-flight and already-resolved

382 + (req_id, key) pairs, so each LookupMsg carries only the keys

383 + first probed in that step. The peer's lookup phase for the id is

384 + still closed by exactly one FetchMsg, which the client contract

385 + guarantees is sent after every lookup for the id has resolved

386 + (see request_blocks / finish). Send-gating is handled by

387 + the injected ``_send`` callback (queues until ConnectAckMsg if

388 + needed).

389 +

390 + Only requests that registered new keys since the last flush are

391 + visited — the ``_flush_pending`` work-list avoids scanning every

392 + live request each scheduler step.

393 + """

394 + for req_id in self._flush_pending:

395 + st = self._requests.get(req_id)

396 + if st is None or not st.unsent:

397 + continue

398 + # Record that the peer now holds lookup state for this id so

399 + # finish knows a terminal empty FetchMsg may be owed.

400 + # Only promote from REGISTERED: once a fetch has gone out

401 + # (FETCH_SENT) a later LookupMsg must not regress the phase, as

402 + # no terminal FetchMsg is owed for an already-fetched request.

403 + if st.phase is ClientPhase.REGISTERED:

404 + st.phase = ClientPhase.PROBING

405 + logger.debug(

406 + "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d",

407 + self._peer_id,

408 + req_id,

409 + len(st.unsent),

410 + )

411 + self._send(

412 + {

413 + TYPE_KEY: LookupMsg.TYPE,

414 + LookupMsg.KV_REQUEST_ID: req_id,

415 + LookupMsg.KEYS: list(st.unsent),

416 + }

417 + )

418 + st.unsent = []

419 + self._flush_pending.clear()

420 +

421 + def on_lookup_resp(

422 + self,

423 + kv_request_id: str,

424 + keys: Sequence[bytes],

425 + hits: Sequence[bool],

426 + ) -> None:

427 + """Apply per-pair hit/miss results from a peer.

428 +

429 + Pairs that don't match a known entry (already cancelled or

430 + never asked) are silently dropped — the producer is free to

431 + split or coalesce responses.

432 + """

433 + n_hit = sum(1 for hit in hits if hit)

434 + logger.debug(

435 + "P2P LOOKUP client %s: RECV LookupRespMsg kv_request_id=%s "

436 + "keys=%d hits=%d misses=%d",

437 + self._peer_id,

438 + kv_request_id,

439 + len(keys),

440 + n_hit,

441 + len(hits) - n_hit,

442 + )

443 + st = self._requests.get(kv_request_id)

444 + if st is None:

445 + return

446 + for h, hit in zip(keys, hits):

447 + key = OffloadKey(h)

448 + if key in st.probes:

449 + st.probes[key] = hit

450 +

451 def collect_results(self) -> list[LoadResult]:

168 - """Walk timeouts and drain completed loads.

452 + """Walk load timeouts and drain completed loads.

453

454 Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg

455 sent and enter the aborting phase. Aborting requests past

456 ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads.

457 +

458 + Lookups have no timeout: an unanswered probe stays None (RETRY)

459 + until finish_request clears it — see ``_ClientRequestState.probes``.

460 """

461 now = time.monotonic()

462 to_remove: list[str] = []

176 - for req_id, req in self._inbound.items():

177 - if req.aborted_at is None:

178 - if now - req.submitted_at >= _LOAD_TIMEOUT_S:

179 - req.aborted_at = now

463 + for req_id in self._active_loads:

464 + st = self._requests[req_id]

465 + assert st.load is not None

466 + load = st.load

467 + if load.aborted_at is None:

468 + if now - load.submitted_at >= _LOAD_TIMEOUT_S:

469 + load.aborted_at = now

470 logger.warning(

471 "P2PSession %s: %s timed out, sending abort",

472 self._peer_id,

@@ -189,11 +479,11 @@ class ClientRole:

479 }

480 )

481 else:

192 - if now - req.aborted_at >= _ABORT_ACK_TIMEOUT_S:

482 + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S:

483 to_remove.append(req_id)

484 self._completed_loads.append(

485 LoadResult(

196 - job_id=req.job_id,

486 + job_id=load.job_id,

487 kv_request_id=req_id,

488 success=False,

489 )

@@ -204,15 +494,31 @@ class ClientRole:

494 req_id,

495 )

496 for req_id in to_remove:

207 - self._inbound.pop(req_id)

497 + self._requests[req_id].load = None

498 + self._active_loads.discard(req_id)

499 + self._maybe_prune(req_id)

500

501 results = self._completed_loads

502 self._completed_loads = []

503 return results

504

213 - def close(self) -> list[tuple[int, str]]:

214 - """Tear down. Returns ``(job_id, kv_request_id)`` for pending loads."""

215 - failed = [(req.job_id, req.kv_request_id) for req in self._inbound.values()]

216 - self._inbound.clear()

505 + def close(self) -> ClientCloseResult:

506 + """Tear down, reporting work the dead peer can no longer complete.

507 +

508 + A request is failed if it has a load in flight (its job fails) or

509 + holds an unresolved symmetric-P2P probe (its lookup() would defer

510 + forever on an answer that can never arrive). See ``ClientCloseResult``.

511 + """

512 + failed_jobs = [

513 + st.load.job_id for st in self._requests.values() if st.load is not None

514 + ]

515 + failed_req_ids = [

516 + req_id

517 + for req_id, st in self._requests.items()

518 + if st.load is not None or any(hit is None for hit in st.probes.values())

519 + ]

520 + self._requests.clear()

521 + self._flush_pending.clear()

522 + self._active_loads.clear()

523 self._completed_loads.clear()

218 - return failed

524 + return ClientCloseResult(failed_jobs=failed_jobs, failed_req_ids=failed_req_ids)

vllm/v1/kv_offload/tiering/p2p/session/protocol.py

+91 -9

@@ -26,6 +26,14 @@ Block Transfer Flow (happy path)

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 + In p2p mode FetchMsg is also the server-side "request finished"

30 + signal for the id: no further ``cb.create_store_job`` will fire

31 + (parked LookupMsg batches are popped, so pending-key resolution

32 + cannot promote a HIT after this point), all server-side lookup

33 + state for the id is released, and ``cb.finish_request`` fires on

34 + each dropped batch. The client emits exactly one FetchMsg per

35 + lookup-touched request, including an empty one when no blocks

36 + end up being fetched.

37 2. Server matches requested blocks against locally stored blocks:

38 - Blocks already available are transferred immediately via RDMA.

39 - Blocks not yet available are recorded as "demanded" and

@@ -119,6 +127,9 @@ class ConnectMsg:

127 BLOCK_LEN: Size in bytes of each block (must match between peers).

128 CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration.

129 Peers with different fingerprints are incompatible.

130 + HASH_SEED: The peer's PYTHONHASHSEED. Block hashes chain from a seed

131 + derived from it, so peers with different values compute different

132 + hashes for identical content and must not exchange blocks.

133 """

134

135 TYPE = "connect"

@@ -128,6 +139,7 @@ class ConnectMsg:

139 NUM_BLOCKS = "num_blocks"

140 BLOCK_LEN = "block_len"

141 CONFIG_FINGERPRINT = "config_fingerprint"

142 + HASH_SEED = "hash_seed"

143

144 @staticmethod

145 def validate(msg: dict) -> None:

@@ -137,6 +149,7 @@ class ConnectMsg:

149 _require_non_neg_int(msg, ConnectMsg.BASE_ADDR)

150 _require_pos_int(msg, ConnectMsg.NUM_BLOCKS)

151 _require_pos_int(msg, ConnectMsg.BLOCK_LEN)

152 + _require(msg, ConnectMsg.HASH_SEED, str)

153

154

155 class ConnectAckMsg:

@@ -165,37 +178,106 @@ class DisconnectMsg:

178

179

180 class FetchMsg:

168 - """Client → Server: request blocks by key.

181 + """Client → Server: request blocks by key and close the lookup phase.

182 +

183 + In p2p mode FetchMsg is also the server-side "request finished"

184 + signal for ``kv_request_id``: on receipt the server (a) fires no

185 + further ``cb.create_store_job`` for this id — parked LookupMsg

186 + batches are popped, so ``_resolve_pending_lookups`` cannot promote

187 + a HIT_PENDING / RETRY key into a fresh pin after this point — and

188 + (b) calls ``cb.finish_request(batch.ctx)`` on each dropped batch

189 + so the TieringManager can release per-batch bookkeeping. In the

190 + all-miss case the client emits an empty FetchMsg (``KEYS``

191 + and ``BLOCK_INDEXES`` both empty) purely to fire this signal.

192

193 Fields:

194 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).

195 + KEYS: List of block keys (OffloadKey bytes). May be empty.

196 + BLOCK_INDEXES: List of remote block indexes (same length as KEYS).

197 """

198

199 TYPE = "fetch"

200 KV_REQUEST_ID = "kv_request_id"

178 - BLOCK_HASHES = "block_hashes"

201 + KEYS = "keys"

202 BLOCK_INDEXES = "block_indexes"

203

204 @staticmethod

205 def validate(msg: dict) -> None:

206 """Raise ValueError if any field has an invalid type or value."""

207 _require(msg, FetchMsg.KV_REQUEST_ID, str)

185 - _require_list(msg, FetchMsg.BLOCK_HASHES)

208 + _require_list(msg, FetchMsg.KEYS)

209 _require_list(msg, FetchMsg.BLOCK_INDEXES)

187 - hashes = msg[FetchMsg.BLOCK_HASHES]

210 + keys = msg[FetchMsg.KEYS]

211 indexes = msg[FetchMsg.BLOCK_INDEXES]

189 - if len(hashes) != len(indexes):

212 + if len(keys) != len(indexes):

213 raise ValueError(

191 - f"block_hashes/block_indexes length mismatch: "

192 - f"{len(hashes)} vs {len(indexes)}"

214 + f"keys/block_indexes length mismatch: {len(keys)} vs {len(indexes)}"

215 )

216 for idx in indexes:

217 if not isinstance(idx, int) or idx < 0:

218 raise ValueError(f"block_indexes: invalid index {idx!r}")

219

220

221 +class LookupMsg:

222 + """Client → Server: probe which block keys the peer holds.

223 +

224 + Sent on the consumer side under symmetric P2P (do_p2p_fetch=true)

225 + after the consumer has aggregated per-block lookups across a

226 + scheduler step. The producer replies with one or more LookupRespMsg

227 + covering the requested keys.

228 +

229 + Fields:

230 + KV_REQUEST_ID: Identifies this lookup transaction.

231 + KEYS: List of block keys (OffloadKey bytes) to probe.

232 + """

233 +

234 + TYPE = "lookup"

235 + KV_REQUEST_ID = "kv_request_id"

236 + KEYS = "keys"

237 +

238 + @staticmethod

239 + def validate(msg: dict) -> None:

240 + """Raise ValueError if any field has an invalid type or value."""

241 + _require(msg, LookupMsg.KV_REQUEST_ID, str)

242 + _require_list(msg, LookupMsg.KEYS)

243 +

244 +

245 +class LookupRespMsg:

246 + """Server → Client: per-key hit/miss answer for a prior LookupMsg.

247 +

248 + Carries two parallel arrays of equal length so each (key,

249 + hit) pair is self-describing. The producer is free to split or

250 + coalesce responses across multiple LookupMsgs for the same

251 + KV_REQUEST_ID — the consumer matches each pair back to its

252 + pending entry by (KV_REQUEST_ID, key).

253 +

254 + Fields:

255 + KV_REQUEST_ID: The lookup transaction this responds to.

256 + KEYS: List of block keys answered by this message.

257 + HITS: Parallel list of bools — True if the producer holds the

258 + corresponding block, False otherwise.

259 + """

260 +

261 + TYPE = "lookup_resp"

262 + KV_REQUEST_ID = "kv_request_id"

263 + KEYS = "keys"

264 + HITS = "hits"

265 +

266 + @staticmethod

267 + def validate(msg: dict) -> None:

268 + """Raise ValueError if any field has an invalid type or value."""

269 + _require(msg, LookupRespMsg.KV_REQUEST_ID, str)

270 + _require_list(msg, LookupRespMsg.KEYS)

271 + _require_list(msg, LookupRespMsg.HITS)

272 + keys = msg[LookupRespMsg.KEYS]

273 + hits = msg[LookupRespMsg.HITS]

274 + if len(keys) != len(hits):

275 + raise ValueError(f"keys/hits length mismatch: {len(keys)} vs {len(hits)}")

276 + for hit in hits:

277 + if not isinstance(hit, bool):

278 + raise ValueError(f"hits: invalid value {hit!r}")

279 +

280 +

281 class TransferDoneMsg:

282 """Server → Client: all blocks transferred for a request.

283

vllm/v1/kv_offload/tiering/p2p/session/server.py

+558 -95

@@ -18,26 +18,32 @@ coordinator's ``_dispatch_message`` can reuse its existing

18 from __future__ import annotations

20 import time

21 -from collections.abc import Callable, Sequence

21 +from collections.abc import Callable, Iterable, Sequence

22 from dataclasses import dataclass, field

23 from typing import TYPE_CHECKING, NamedTuple

25 from vllm.logger import init_logger

26 -from vllm.v1.kv_offload.base import OffloadKey

26 +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext

27 from vllm.v1.kv_offload.tiering.p2p.session.protocol import (

28 TYPE_KEY,

29 AbortAckMsg,

30 + LookupRespMsg,

31 TransferDoneMsg,

32 )

34 if TYPE_CHECKING:

34 - from vllm.v1.kv_offload.tiering.base import JobId

35 + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager

36 from vllm.v1.kv_offload.tiering.p2p.data import DataTransport

38 logger = init_logger(__name__)

40 _STORE_TIMEOUT_S = 30.0

41 _CANCEL_DRAIN_TIMEOUT_S = 10.0

42 +# Cap on time the server holds a HIT_PENDING / RETRY key from an inbound

43 +# LookupMsg before falling back to MISS. Long enough that an in-flight

44 +# primary write or a just-started promotion typically completes; short

45 +# enough that the consumer doesn't sit idle on a stuck producer.

46 +_LOOKUP_PENDING_TIMEOUT_S = 5.0

49 class StoreResult(NamedTuple):

@@ -69,8 +75,8 @@ class _MatchResult(NamedTuple):

75 class _OutboundRequestState:

76 """Server-role state for a single peer fetch request.

72 - The owning ``kv_request_id`` is the dict key in

73 - ``ServerRole._outbound`` and is not duplicated on the value.

78 + The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key

79 + and is not duplicated on the value.

80 """

82 demand_received: bool = False

@@ -90,7 +96,7 @@ class _OutboundRequestState:

97 def add_stored_blocks(

98 self,

93 - block_hashes: Sequence[OffloadKey],

99 + keys: Sequence[OffloadKey],

100 block_ids: Sequence[int],

101 job_id: int,

102 ) -> _MatchResult:

@@ -98,13 +104,13 @@ class _OutboundRequestState:

104 self.pending_job_ids.add(job_id)

105 local_idxs: list[int] = []

106 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)

107 + for key, local_idx in zip(keys, block_ids):

108 + remote_idx = self.demanded.pop(key, None)

109 if remote_idx is not None:

110 local_idxs.append(local_idx)

111 remote_idxs.append(remote_idx)

112 else:

107 - self.available[block_hash] = (job_id, local_idx)

113 + self.available[key] = (job_id, local_idx)

114 return _MatchResult(

115 local_idxs=local_idxs,

116 remote_idxs=remote_idxs,

@@ -113,25 +119,25 @@ class _OutboundRequestState:

119

120 def add_fetch_demand(

121 self,

116 - block_hashes: Sequence[OffloadKey],

122 + keys: Sequence[OffloadKey],

123 block_indexes: Sequence[int],

124 ) -> _MatchResult:

125 """Register the peer's fetch demand. Returns matched pairs."""

126 self.demand_received = True

121 - self.remaining = len(block_hashes)

127 + self.remaining = len(keys)

128

129 local_idxs: list[int] = []

130 remote_idxs: list[int] = []

131 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)

132 + for key, remote_idx in zip(keys, block_indexes):

133 + stored_entry = self.available.pop(key, None)

134 if stored_entry is not None:

135 stored_job_id, local_idx = stored_entry

136 local_idxs.append(local_idx)

137 remote_idxs.append(remote_idx)

138 job_ids.add(stored_job_id)

139 else:

134 - self.demanded[block_hash] = remote_idx

140 + self.demanded[key] = remote_idx

141 return _MatchResult(

142 local_idxs=local_idxs,

143 remote_idxs=remote_idxs,

@@ -139,6 +145,78 @@ class _OutboundRequestState:

145 )

146

147

148 +@dataclass

149 +class _ActiveLookup:

150 + """In-flight state for one inbound LookupMsg.

151 +

152 + Aggregates per-key HIT/MISS resolutions and defers the single

153 + outbound LookupRespMsg until every key has been resolved or the

154 + ``deadline`` fires (remaining ``pending`` keys then force-MISS).

155 + Exactly one LookupRespMsg is emitted per LookupMsg, carrying every

156 + key in the original wire order.

157 + """

158 +

159 + lookup_id: int

160 + kv_request_id: str

161 + ctx: ReqContext

162 + # Keys from the inbound LookupMsg, preserved in wire order so

163 + # the aggregated response goes back in the same order.

164 + keys: list[OffloadKey] = field(default_factory=list)

165 + # Per-key resolution: True = HIT, False = MISS. A key is present

166 + # here once definitively resolved; still-pending keys are only

167 + # in ``pending``.

168 + resolved: dict[OffloadKey, bool] = field(default_factory=dict)

169 + # Keys still awaiting resolution (HIT_PENDING / RETRY from

170 + # ``parent.lookup``); re-polled by ``_resolve_pending_lookups``.

171 + pending: set[OffloadKey] = field(default_factory=set)

172 + # Absolute deadline (``time.monotonic``). Once reached, remaining

173 + # ``pending`` keys are force-resolved to MISS so the consumer

174 + # can fall back instead of waiting on a stuck producer.

175 + deadline: float = 0.0

176 +

177 +

178 +class _PendingLookup(NamedTuple):

179 + """A raw inbound LookupMsg awaiting resolution.

180 +

181 + Enqueued by ``on_lookup`` during dispatch and drained by the next

182 + ``serve_external_requests``. The deadline for any resulting

183 + HIT_PENDING / RETRY key is measured from ``enqueued_at``.

184 + """

185 +

186 + keys: list[OffloadKey]

187 + enqueued_at: float

188 +

189 +

190 +@dataclass

191 +class _ServerRequestState:

192 + """Per-kv_request_id server-side state.

193 +

194 + Consolidates the outbound serve/transfer state, the symmetric-P2P

195 + inbound-lookup state, abort bookkeeping, and the inflight-transfer

196 + count for one kv_request_id. The owning id is the ``_requests`` dict

197 + key and is not duplicated here. An entry is dropped once every field

198 + is idle — see ``ServerRole._maybe_prune``.

199 + """

200 +

201 + # Outbound serve state (PD + symmetric producer side). None until the

202 + # first add_stored_blocks / on_fetch; reset to None on finalize/abort.

203 + outbound: _OutboundRequestState | None = None

204 + # Raw inbound LookupMsgs not yet processed against the ParentManager.

205 + pending_lookups: list[_PendingLookup] = field(default_factory=list)

206 + # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by

207 + # the (globally unique) lookup_id and re-polled each serve.

208 + lookups: dict[int, _ActiveLookup] = field(default_factory=dict)

209 + # Start time of a pending abort drain (``time.monotonic``); None when

210 + # no abort is in progress.

211 + abort_started_at: float | None = None

212 + # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync

213 + # via _inflight_add / _inflight_pop so a non-empty set is an exact

214 + # "has any inflight transfer" predicate and the abort drain can

215 + # enumerate this request's transfers without scanning all of

216 + # ``_inflight``.

217 + inflight_tids: set[int] = field(default_factory=set)

218 +

219 +

220 class ServerRole:

221 """Server-side store/serve state machine for one peer session.

222

@@ -157,22 +235,55 @@ class ServerRole:

235 self._transport = transport

236 self._send = send

237

160 - self._outbound: dict[str, _OutboundRequestState] = {}

238 + # All per-kv_request_id state lives here. Entries are created

239 + # lazily and dropped by _maybe_prune once every field is idle.

240 + self._requests: dict[str, _ServerRequestState] = {}

241 + # kv_request_ids with lookup work (unprocessed pending_lookups or

242 + # parked lookups) for the next serve to visit — the work-list that

243 + # keeps serve_external_requests from scanning every request.

244 + self._serve_pending: set[str] = set()

245 # transfer_id → xfer. Mutate ONLY via _inflight_add / _inflight_pop

162 - # so the per-request count below stays in sync.

246 + # so the per-request inflight_tids stays in sync.

247 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] = {}

248 self._store_jobs: dict[int, float] = {} # job_id → submitted_at

171 - self._pending_aborts: dict[str, float] = {} # kv_request_id → start

249 # StoreResults queued by _finalize_outbound for the next poll

250 # tick to surface. Mirrors the deferred-result pattern used for

251 # load timeouts.

252 self._pending_store_results: list[StoreResult] = []

253 + # Synthetic lookup ctxs whose ``on_request_finished`` still needs to

254 + # fire but which were closed outside a serve window (FetchMsg / local

255 + # finish popped their parked lookup). Drained via

256 + # ``parent.on_request_finished`` in ``serve_external_requests``.

257 + self._finished_lookup_ctxs: list[ReqContext] = []

258 + self._lookup_id_counter: int = 0

259 + # kv_request_ids with a parked abort awaiting drain — work-list so

260 + # drain_pending_aborts doesn't scan every request each poll tick.

261 + self._parked_aborts: set[str] = set()

262 +

263 + # ------------------------------------------------------------------

264 + # State helpers

265 + # ------------------------------------------------------------------

266 +

267 + def _get_or_create_request(self, kv_request_id: str) -> _ServerRequestState:

268 + """Get or create the state entry for a kv_request_id."""

269 + st = self._requests.get(kv_request_id)

270 + if st is None:

271 + st = _ServerRequestState()

272 + self._requests[kv_request_id] = st

273 + return st

274 +

275 + def _maybe_prune(self, kv_request_id: str) -> None:

276 + """Drop the entry once it holds no live state."""

277 + st = self._requests.get(kv_request_id)

278 + if (

279 + st is not None

280 + and st.outbound is None

281 + and not st.inflight_tids

282 + and not st.lookups

283 + and not st.pending_lookups

284 + and st.abort_started_at is None

285 + ):

286 + del self._requests[kv_request_id]

287

288 # ------------------------------------------------------------------

289 # Public API

@@ -187,19 +298,41 @@ class ServerRole:

298 ) -> None:

299 """New blocks stored locally — match against pending fetch demand."""

300 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:

301 + st = self._get_or_create_request(kv_request_id)

302 + if st.outbound is None:

303 + st.outbound = _OutboundRequestState()

304 + result = st.outbound.add_stored_blocks(keys, block_ids, job_id)

305 + if result.local_idxs and st.outbound.demand_received:

306 self._submit_transfer(kv_request_id, result)

307

308 def on_fetch(

309 self,

310 kv_request_id: str,

198 - block_hashes: Sequence[OffloadKey],

311 + keys: Sequence[OffloadKey],

312 block_indexes: Sequence[int],

313 ) -> None:

314 """Handle a FetchMsg from the peer.

315

316 + In p2p mode FetchMsg is the server-side "request finished"

317 + signal for ``kv_request_id``. Three consequences flow from that:

318 +

319 + - No further ``parent.create_store_job`` will fire for this id:

320 + the request's parked ``lookups`` are popped here (via

321 + ``_finish_inbound_lookups``) before the next

322 + ``serve_external_requests`` runs ``_resolve_pending_lookups``,

323 + so any HIT_PENDING / RETRY key that would otherwise later

324 + promote to HIT and pin a slot is dropped instead. Any raw

325 + not-yet-processed LookupMsg for this id is dropped too.

326 + - All server-side lookup state for the id is cleaned up (the

327 + ``lookups`` entries themselves).

328 + - The synthetic ctx is queued for ``parent.on_request_finished``

329 + (fired by the next ``serve_external_requests``) so the

330 + TieringManager can release per-lookup bookkeeping.

331 +

332 + The client contract guarantees exactly one FetchMsg per

333 + lookup-touched request — including an empty one when no

334 + blocks end up being fetched.

335 +

336 Raises ``ValueError`` on a duplicate fetch for the same

337 ``kv_request_id``; the coordinator's dispatch loop turns that

338 into a protocol-error disconnect.

@@ -208,18 +341,28 @@ class ServerRole:

341 "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d",

342 self._peer_id,

343 kv_request_id,

211 - len(block_hashes),

344 + len(keys),

345 )

213 - existing = self._outbound.get(kv_request_id)

346 + st = self._requests.get(kv_request_id)

347 + existing = st.outbound if st is not None else None

348 if existing is not None and existing.demand_received:

349 # A second fetch for the same kv_request_id would overwrite

350 # `remaining` and leak inflight bookkeeping. Treat as a

351 # protocol violation.

352 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)

353 + st = self._get_or_create_request(kv_request_id)

354 + if st.outbound is None:

355 + st.outbound = _OutboundRequestState()

356 + req = st.outbound

357 + result = req.add_fetch_demand(keys, block_indexes)

358 if result.local_idxs:

359 self._submit_transfer(kv_request_id, result)

360 + # Close the peer's request as far as the server's lookup phase

361 + # is concerned: pop parked lookups so no further

362 + # ``parent.create_store_job`` fires for this id, and queue their

363 + # ctxs for ``parent.on_request_finished``. Done before the

364 + # finalize path below so bookkeeping releases in-order.

365 + self._finish_inbound_lookups(kv_request_id)

366 # Prefiller-first mode: finish_request may have run before

367 # fetch arrived. If so, finalize once we know what was

368 # demanded — fully satisfied → success, else early-fail.

@@ -230,9 +373,9 @@ class ServerRole:

373 """Handle an AbortFetchMsg from the peer."""

374 # Abort for an unknown id may be a benign race/duplicate or a

375 # 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 - ):

376 + st = self._requests.get(kv_request_id)

377 + has_outbound = st is not None and st.outbound is not None

378 + if not has_outbound and not self._has_inflight_for(kv_request_id):

379 logger.warning(

380 "P2PSession %s: abort_fetch for unknown kv_request_id=%s "

381 "(no outbound or inflight state); benign race or stale",

@@ -242,9 +385,285 @@ class ServerRole:

385 # Idempotent: receiving AbortFetchMsg again before we've sent the

386 # ack just triggers another drain attempt without resetting the

387 # deadline.

245 - self._pending_aborts.setdefault(kv_request_id, time.monotonic())

388 + st = self._get_or_create_request(kv_request_id)

389 + if st.abort_started_at is None:

390 + st.abort_started_at = time.monotonic()

391 + self._parked_aborts.add(kv_request_id)

392 self._drain_abort(kv_request_id)

393

394 + def on_lookup(

395 + self,

396 + kv_request_id: str,

397 + keys: Sequence[OffloadKey],

398 + ) -> None:

399 + """Enqueue a LookupMsg from a symmetric-P2P consumer.

400 +

401 + Dispatch runs during ``session.poll()`` where the

402 + :class:`ParentManager` handle is not available, so this only

403 + records the raw request. It is resolved — querying the tiering

404 + manager and emitting the aggregated ``LookupRespMsg`` — by the

405 + next :meth:`serve_external_requests`, the sole window in which

406 + parent calls are valid.

407 + """

408 + logger.debug(

409 + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d",

410 + self._peer_id,

411 + kv_request_id,

412 + len(keys),

413 + )

414 + self._get_or_create_request(kv_request_id).pending_lookups.append(

415 + _PendingLookup(keys=list(keys), enqueued_at=time.monotonic())

416 + )

417 + self._serve_pending.add(kv_request_id)

418 +

419 + def serve_external_requests(self, parent: ParentManager) -> None:

420 + """Resolve inbound peer lookups against the tiering manager.

421 +

422 + Called once per scheduler step with a ``parent`` handle valid

423 + only for this call. Drains newly-enqueued LookupMsgs, re-polls

424 + any parked HIT_PENDING / RETRY keys, and releases the

425 + bookkeeping for lookups closed since the last serve.

426 + """

427 + for kv_request_id in list(self._serve_pending):

428 + st = self._requests.get(kv_request_id)

429 + if st is None:

430 + self._serve_pending.discard(kv_request_id)

431 + continue

432 + if st.pending_lookups:

433 + pending = st.pending_lookups

434 + st.pending_lookups = []

435 + for pl in pending:

436 + self._process_inbound_lookup(

437 + kv_request_id, pl.keys, pl.enqueued_at, parent

438 + )

439 + self._resolve_pending_lookups(kv_request_id, parent)

440 + st = self._requests.get(kv_request_id)

441 + if st is None or (not st.pending_lookups and not st.lookups):

442 + self._serve_pending.discard(kv_request_id)

443 + self._maybe_prune(kv_request_id)

444 +

445 + if self._finished_lookup_ctxs:

446 + for ctx in self._finished_lookup_ctxs:

447 + parent.on_request_finished(ctx)

448 + self._finished_lookup_ctxs = []

449 +

450 + def _poll_lookup_keys(

451 + self,

452 + lookup: _ActiveLookup,

453 + keys: Iterable[OffloadKey],

454 + parent: ParentManager,

455 + ) -> list[OffloadKey]:

456 + """Poll ``keys`` against the tiering manager and pin any HITs.

457 +

458 + For each key not already definitively resolved, query

459 + ``parent.lookup`` and record the outcome on ``lookup``: HIT / MISS

460 + land in ``resolved`` and clear ``pending``; HIT_PENDING / RETRY park

461 + in ``pending`` for a later serve. Newly-HIT keys are pinned in one

462 + batch via :meth:`_pin_and_register_hits` and also returned (for

463 + caller logging).

464 +

465 + Shared by the first-sighting pass (:meth:`_process_inbound_lookup`)

466 + and the re-poll pass (:meth:`_resolve_pending_lookups`); callers pass

467 + a de-duplicated ``keys`` collection.

468 + """

469 + new_hits: list[OffloadKey] = []

470 + for h in keys:

471 + if h in lookup.resolved:

472 + continue

473 + result = parent.lookup(h, lookup.ctx)

474 + if result is LookupResult.HIT:

475 + new_hits.append(h)

476 + lookup.resolved[h] = True

477 + lookup.pending.discard(h)

478 + elif result is LookupResult.MISS:

479 + lookup.resolved[h] = False

480 + lookup.pending.discard(h)

481 + else:

482 + lookup.pending.add(h)

483 + if new_hits:

484 + self._pin_and_register_hits(

485 + lookup.kv_request_id, new_hits, lookup.ctx, parent

486 + )

487 + return new_hits

488 +

489 + def _process_inbound_lookup(

490 + self,

491 + kv_request_id: str,

492 + keys: list[OffloadKey],

493 + enqueued_at: float,

494 + parent: ParentManager,

495 + ) -> None:

496 + """Resolve one enqueued LookupMsg against ``parent``.

497 +

498 + For each key, query the tiering manager via ``parent.lookup``

499 + and pin any HITs immediately via ``parent.create_store_job``

500 + (plumbed into the existing ``add_stored_blocks`` matching path so

501 + the eventual FetchMsg finds them). HIT_PENDING / RETRY keys park

502 + in :class:`_ActiveLookup` for re-polling by

503 + :meth:`_resolve_pending_lookups`.

504 +

505 + The outbound ``LookupRespMsg`` is deferred until every key has

506 + settled to HIT or MISS (or the batch-level ``deadline`` fires,

507 + forcing any stragglers to MISS). One LookupRespMsg goes out per

508 + LookupMsg — carrying every key in wire order — after which

509 + ``parent.on_request_finished`` fires and the entry is dropped.

510 + """

511 + self._lookup_id_counter += 1

512 + lookup_id = self._lookup_id_counter

513 + ctx = ReqContext(req_id=f"p2p:{self._peer_id}:{kv_request_id}:lu{lookup_id}")

514 + lookup = _ActiveLookup(

515 + lookup_id=lookup_id,

516 + kv_request_id=kv_request_id,

517 + ctx=ctx,

518 + keys=list(keys),

519 + deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S,

520 + )

521 +

522 + # Open per-request bookkeeping for this synthetic ctx before the

523 + # first lookup; released by ``on_request_finished`` once every

524 + # key has settled.

525 + parent.on_new_request(ctx)

526 +

527 + # dict.fromkeys de-duplicates keys within the LookupMsg while

528 + # preserving wire order — each unique key is polled once.

529 + hit_keys = self._poll_lookup_keys(lookup, dict.fromkeys(lookup.keys), parent)

530 +

531 + logger.debug(

532 + "P2P LOOKUP server %s: RESOLVED kv_request_id=%s hits=%d misses=%d "

533 + "pending=%d",

534 + self._peer_id,

535 + kv_request_id,

536 + len(hit_keys),

537 + sum(1 for v in lookup.resolved.values() if not v),

538 + len(lookup.pending),

539 + )

540 +

541 + if lookup.pending:

542 + self._get_or_create_request(kv_request_id).lookups[lookup_id] = lookup

543 + else:

544 + # Every key resolved on first sight — emit the aggregated

545 + # response now and close the synthetic request.

546 + self._finalize_lookup(lookup, parent)

547 +

548 + def _pin_and_register_hits(

549 + self,

550 + kv_request_id: str,

551 + keys: list[OffloadKey],

552 + ctx: ReqContext,

553 + parent: ParentManager,

554 + ) -> None:

555 + """Pin primary slots for HIT keys and feed them into the

556 + existing ``add_stored_blocks`` matching path.

557 +

558 + Caller has already confirmed every key is HIT (single-threaded

559 + scheduler ⇒ no eviction race), so the JobMetadata returned by

560 + ``parent.create_store_job`` carries parallel ``keys``/``block_ids``

561 + of length ``len(keys)``.

562 + """

563 + meta = parent.create_store_job(keys, ctx)

564 + self.add_stored_blocks(

565 + kv_request_id,

566 + list(meta.keys),

567 + list(meta.block_ids),

568 + meta.job_id,

569 + )

570 +

571 + def _resolve_pending_lookups(

572 + self, kv_request_id: str, parent: ParentManager

573 + ) -> None:

574 + """Re-poll a request's deferred LookupMsg keys; finalize when ready.

575 +

576 + Walks every parked :class:`_ActiveLookup` for ``kv_request_id`` and

577 + re-calls ``parent.lookup`` per still-pending key, moving HIT/MISS

578 + results into ``resolved``. If ``deadline`` has passed, remaining

579 + ``pending`` keys are force-resolved to MISS so the consumer

580 + can fall back instead of waiting on a stuck producer. Newly-HIT

581 + keys are pinned via ``parent.create_store_job`` in one call per

582 + affected lookup. Lookups whose ``pending`` empties out get their

583 + aggregated LookupRespMsg sent by :meth:`_finalize_lookup`.

584 + """

585 + st = self._requests.get(kv_request_id)

586 + if st is None or not st.lookups:

587 + return

588 + now = time.monotonic()

589 + finished_lookups: list[int] = []

590 + for lookup_id, lookup in st.lookups.items():

591 + self._poll_lookup_keys(lookup, list(lookup.pending), parent)

592 +

593 + if lookup.pending and now >= lookup.deadline:

594 + for h in lookup.pending:

595 + lookup.resolved[h] = False

596 + lookup.pending.clear()

597 +

598 + if not lookup.pending:

599 + finished_lookups.append(lookup_id)

600 +

601 + for lookup_id in finished_lookups:

602 + lookup = st.lookups.pop(lookup_id)

603 + self._finalize_lookup(lookup, parent)

604 +

605 + def _finalize_lookup(self, lookup: _ActiveLookup, parent: ParentManager) -> None:

606 + """Emit the aggregated LookupRespMsg and close the synthetic request.

607 +

608 + Called once per lookup when ``pending`` is empty — either every

609 + key resolved to HIT or MISS, or the deadline forced remaining

610 + stragglers to MISS. Preserves the wire order of the inbound

611 + LookupMsg so the client can zip keys and hits positionally.

612 + """

613 + if lookup.keys:

614 + hits = [lookup.resolved[h] for h in lookup.keys]

615 + n_hit = sum(1 for v in hits if v)

616 + logger.debug(

617 + "P2P LOOKUP server %s: SEND LookupRespMsg kv_request_id=%s "

618 + "keys=%d hits=%d misses=%d",

619 + self._peer_id,

620 + lookup.kv_request_id,

621 + len(lookup.keys),

622 + n_hit,

623 + len(hits) - n_hit,

624 + )

625 + self._send(

626 + {

627 + TYPE_KEY: LookupRespMsg.TYPE,

628 + LookupRespMsg.KV_REQUEST_ID: lookup.kv_request_id,

629 + LookupRespMsg.KEYS: list(lookup.keys),

630 + LookupRespMsg.HITS: hits,

631 + }

632 + )

633 + parent.on_request_finished(lookup.ctx)

634 +

635 + def _finish_inbound_lookups(self, kv_request_id: str) -> None:

636 + """Close the server-side lookup phase for ``kv_request_id``.

637 +

638 + Pops every parked ``_ActiveLookup`` for this id (so

639 + ``_resolve_pending_lookups`` cannot promote a HIT_PENDING /

640 + RETRY key into a fresh ``parent.create_store_job`` after this

641 + point) and queues each ``lookup.ctx`` for

642 + ``parent.on_request_finished`` (fired by the next

643 + ``serve_external_requests``, since no parent handle is available

644 + during dispatch) so the TieringManager can release per-lookup

645 + bookkeeping. Any still-unprocessed raw LookupMsg for this id is

646 + dropped — it never got ``on_new_request``, so nothing is owed.

647 + The aggregated LookupRespMsg is skipped — the client already

648 + knows the request is over (it just sent a terminal FetchMsg, or

649 + is finishing locally).

650 +

651 + Called on the two events that mean "no more lookup traffic for

652 + ``kv_request_id`` is expected on this session": the terminal

653 + FetchMsg from the peer (client contract: exactly one FetchMsg

654 + per lookup-touched request, even if empty), and a local

655 + ``finish``. Whichever fires second is a no-op.

656 + """

657 + st = self._requests.get(kv_request_id)

658 + if st is None:

659 + return

660 + st.pending_lookups.clear()

661 + for lookup in st.lookups.values():

662 + self._finished_lookup_ctxs.append(lookup.ctx)

663 + st.lookups.clear()

664 + self._serve_pending.discard(kv_request_id)

665 + self._maybe_prune(kv_request_id)

666 +

667 def finish(self, kv_request_id: str) -> None:

668 """Mark an outbound request finishing.

669

@@ -253,13 +672,23 @@ class ServerRole:

672 peer to stop waiting (TransferDoneMsg success=False) instead of

673 letting it hit _LOAD_TIMEOUT_S.

674

675 + Also drops any in-flight lookups for this kv_request_id and

676 + queues their ctxs for ``parent.on_request_finished`` so the

677 + TieringManager can release per-request bookkeeping. (For

678 + symmetric P2P this path is rarely hit since the producer has no

679 + local request lifecycle for the consumer's id; this cleanup is

680 + mostly active on the PD side.)

681 +

682 If the decoder hasn't sent fetch yet (no demand received),

683 defer — on_fetch will finalize once demand arrives.

684

685 If inflight transfers exist for this id, defer — the last

686 completing transfer in collect_results will fire the message.

687 """

262 - req = self._outbound.get(kv_request_id)

688 + self._finish_inbound_lookups(kv_request_id)

689 +

690 + st = self._requests.get(kv_request_id)

691 + req = st.outbound if st is not None else None

692 if req is None:

693 return

694 req.finishing = True

@@ -268,14 +697,19 @@ class ServerRole:

697 if self._has_inflight_for(kv_request_id):

698 return

699 # Remaining > 0 here: if it had hit 0, the poll-done success

271 - # branch would have already popped _outbound and we'd have

700 + # branch would have already cleared outbound and we'd have

701 # returned at `req is None` above. Helper derives success from

702 # remaining and emits StoreResult(success=False) for any

703 # leftover pending jobs.

704 self._finalize_outbound(kv_request_id)

705

706 def collect_results(self) -> list[StoreResult]:

278 - """Drain timeouts, deferred results, and transport completions."""

707 + """Drain timeouts, deferred results, and transport completions.

708 +

709 + Inbound LookupMsg resolution (including re-polling HIT_PENDING /

710 + RETRY keys) is NOT done here — it runs in

711 + ``serve_external_requests`` where the ParentManager is available.

712 + """

713 results: list[StoreResult] = self._timeout_pending_store_jobs()

714

715 if self._pending_store_results:

@@ -306,15 +740,9 @@ class ServerRole:

740 tid,

741 )

742 continue

309 - req = self._outbound.get(xfer.kv_request_id)

310 - for job_id in xfer.job_ids:

311 - if self._store_jobs.pop(job_id, None) is None:

312 - # Already reported (timeout, cancellation, etc.) —

313 - # don't double-emit a contradictory success result.

314 - continue

315 - results.append(StoreResult(job_id=job_id, success=True))

316 - if req is not None:

317 - req.pending_job_ids.discard(job_id)

743 + results.extend(self._settle_xfer_jobs(xfer, success=True))

744 + st = self._requests.get(xfer.kv_request_id)

745 + req = st.outbound if st is not None else None

746 if req is not None and req.demand_received:

747 req.remaining -= xfer.block_count

748 assert req.remaining >= 0, (

@@ -324,6 +752,7 @@ class ServerRole:

752 self._finalize_outbound(xfer.kv_request_id, success=True)

753 elif req.finishing and not self._has_inflight_for(xfer.kv_request_id):

754 self._finalize_outbound(xfer.kv_request_id, success=False)

755 + self._maybe_prune(xfer.kv_request_id)

756

757 failed_kv_request_ids: set[str] | None = None

758 for tid in poll_result.failed:

@@ -341,16 +770,11 @@ class ServerRole:

770 if failed_kv_request_ids is None:

771 failed_kv_request_ids = set()

772 failed_kv_request_ids.add(xfer.kv_request_id)

344 - req_for_xfer = self._outbound.get(xfer.kv_request_id)

345 - for job_id in xfer.job_ids:

346 - if self._store_jobs.pop(job_id, None) is None:

347 - # Already reported (timeout, cancellation, etc.) —

348 - # don't double-emit.

349 - continue

350 - results.append(StoreResult(job_id=job_id, success=False))

351 - if req_for_xfer is not None:

352 - req_for_xfer.pending_job_ids.discard(job_id)

353 - req = self._outbound.pop(xfer.kv_request_id, None)

773 + results.extend(self._settle_xfer_jobs(xfer, success=False))

774 + st = self._requests.get(xfer.kv_request_id)

775 + req = st.outbound if st is not None else None

776 + if st is not None:

777 + st.outbound = None

778 if req is not None and req.demand_received:

779 self._send(

780 {

@@ -359,6 +783,7 @@ class ServerRole:

783 TransferDoneMsg.SUCCESS: False,

784 }

785 )

786 + self._maybe_prune(xfer.kv_request_id)

787

788 # Cancel other inflight for the same failed kv_request_ids

789 if failed_kv_request_ids:

@@ -370,6 +795,8 @@ class ServerRole:

795 for tid in ids_to_cancel:

796 self._inflight_pop(tid)

797 self._transport.cancel(ids_to_cancel)

798 + for kv_request_id in failed_kv_request_ids:

799 + self._maybe_prune(kv_request_id)

800

801 return results

802

@@ -384,54 +811,90 @@ class ServerRole:

811

812 def drain_pending_aborts(self) -> None:

813 """Re-attempt every parked abort once per poll tick."""

387 - if not self._pending_aborts:

388 - return

389 - for kv_request_id in list(self._pending_aborts):

814 + for kv_request_id in list(self._parked_aborts):

815 self._drain_abort(kv_request_id)

816

392 - def close(self) -> list[int]:

393 - """Tear down. Cancels inflight, returns failed store job ids."""

817 + def close(self) -> tuple[list[int], list[ReqContext]]:

818 + """Tear down. Cancels inflight.

819 +

820 + Returns ``(failed_store_job_ids, failed_serves)`` where

821 + ``failed_serves`` are synthetic lookup ctxs still owing a

822 + ``parent.on_request_finished``. The session is going away with no

823 + parent handle in hand, so the manager flushes these in its next

824 + ``serve_external_requests``.

825 + """

826 failed_stores = list(self._store_jobs.keys())

827 self._store_jobs.clear()

828 if self._inflight:

829 self._transport.cancel(list(self._inflight.keys()))

830 self._inflight.clear()

399 - self._inflight_per_req.clear()

400 - self._outbound.clear()

401 - self._pending_aborts.clear()

831 self._pending_store_results.clear()

403 - return failed_stores

832 + # Surface every synthetic ctx still owing on_request_finished so

833 + # the manager can release the TieringManager's per-request

834 + # bookkeeping: parked lookups plus any already queued from a

835 + # FetchMsg / finish that closed them before this teardown.

836 + failed_serves = [

837 + lu.ctx for st in self._requests.values() for lu in st.lookups.values()

838 + ]

839 + failed_serves.extend(self._finished_lookup_ctxs)

840 + self._requests.clear()

841 + self._serve_pending.clear()

842 + self._parked_aborts.clear()

843 + self._finished_lookup_ctxs.clear()

844 + return failed_stores, failed_serves

845 +

846 + @property

847 + def has_inflight_transfers(self) -> bool:

848 + """True if any outbound store transfer is still in flight."""

849 + return bool(self._inflight)

850

851 # ------------------------------------------------------------------

852 # Internal — inflight bookkeeping

853 # ------------------------------------------------------------------

854

855 def _has_inflight_for(self, kv_request_id: str) -> bool:

410 - return kv_request_id in self._inflight_per_req

856 + st = self._requests.get(kv_request_id)

857 + return st is not None and bool(st.inflight_tids)

858

859 def _inflight_add(self, tid: int, xfer: _InflightXfer) -> None:

413 - """Insert an inflight transfer and bump the per-request count."""

860 + """Insert an inflight transfer and record it on the request."""

861 self._inflight[tid] = xfer

415 - self._inflight_per_req[xfer.kv_request_id] = (

416 - self._inflight_per_req.get(xfer.kv_request_id, 0) + 1

417 - )

862 + self._get_or_create_request(xfer.kv_request_id).inflight_tids.add(tid)

863

864 def _inflight_pop(self, tid: int) -> _InflightXfer | None:

420 - """Pop an inflight transfer and decrement the per-request count.

865 + """Pop an inflight transfer and drop it from the request's set.

866

422 - Removes the per-request entry once the count hits zero so the

423 - dict stays bounded and `_has_inflight_for` remains exact.

867 + Callers are responsible for the ``_maybe_prune`` that may follow once

868 + the request's other state has also cleared.

869 """

870 xfer = self._inflight.pop(tid, None)

871 if xfer is None:

872 return None

428 - new_count = self._inflight_per_req.get(xfer.kv_request_id, 0) - 1

429 - if new_count > 0:

430 - self._inflight_per_req[xfer.kv_request_id] = new_count

431 - else:

432 - self._inflight_per_req.pop(xfer.kv_request_id, None)

873 + st = self._requests.get(xfer.kv_request_id)

874 + if st is not None:

875 + st.inflight_tids.discard(tid)

876 return xfer

877

878 + def _settle_xfer_jobs(

879 + self, xfer: _InflightXfer, success: bool

880 + ) -> list[StoreResult]:

881 + """Emit StoreResults for a completed transfer's store jobs.

882 +

883 + Pops each attached job from ``_store_jobs`` and clears it from the

884 + request's pending set. A job already popped (via timeout, cancel,

885 + etc.) is skipped so we never double-emit a contradictory result.

886 + """

887 + results: list[StoreResult] = []

888 + st = self._requests.get(xfer.kv_request_id)

889 + req = st.outbound if st is not None else None

890 + for job_id in xfer.job_ids:

891 + if self._store_jobs.pop(job_id, None) is None:

892 + continue

893 + results.append(StoreResult(job_id=job_id, success=success))

894 + if req is not None:

895 + req.pending_job_ids.discard(job_id)

896 + return results

897 +

898 # ------------------------------------------------------------------

899 # Internal — finalize / abort drain

900 # ------------------------------------------------------------------

@@ -452,9 +915,10 @@ class ServerRole:

915 The same flag is used for both the peer's TransferDoneMsg and

916 the StoreResult(s) emitted for any leftover pending job_ids.

917 """

455 - req = self._outbound.pop(kv_request_id, None)

456 - if req is None:

457 - return

918 + st = self._requests[kv_request_id]

919 + assert st.outbound is not None

920 + req = st.outbound

921 + st.outbound = None

922 if success is None:

923 success = req.remaining == 0

924 for job_id in req.pending_job_ids:

@@ -469,6 +933,7 @@ class ServerRole:

933 TransferDoneMsg.SUCCESS: success,

934 }

935 )

936 + self._maybe_prune(kv_request_id)

937

938 def _drain_abort(self, kv_request_id: str) -> None:

939 """One drain attempt for a pending abort.

@@ -479,21 +944,15 @@ class ServerRole:

944 inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to

945 ``mode="immediate"`` and acks anyway.

946 """

482 - self._outbound.pop(kv_request_id, None)

483 - ids = [

484 - tid

485 - for tid, xfer in self._inflight.items()

486 - if xfer.kv_request_id == kv_request_id

487 - ]

947 + st = self._requests[kv_request_id]

948 + st.outbound = None

949 + ids = list(st.inflight_tids)

950 if not ids:

951 self._finalize_abort(kv_request_id)

952 return

953

492 - started_at = self._pending_aborts.get(kv_request_id)

493 - expired = (

494 - started_at is not None

495 - and time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S

496 - )

954 + assert st.abort_started_at is not None

955 + expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S

956 if expired:

957 for tid in ids:

958 self._inflight_pop(tid)

@@ -521,13 +980,16 @@ class ServerRole:

980 self._finalize_abort(kv_request_id)

981

982 def _finalize_abort(self, kv_request_id: str) -> None:

524 - self._pending_aborts.pop(kv_request_id, None)

983 + st = self._requests[kv_request_id]

984 + st.abort_started_at = None

985 + self._parked_aborts.discard(kv_request_id)

986 self._send(

987 {

988 TYPE_KEY: AbortAckMsg.TYPE,

989 AbortAckMsg.KV_REQUEST_ID: kv_request_id,

990 }

991 )

992 + self._maybe_prune(kv_request_id)

993

994 # ------------------------------------------------------------------

995 # Internal — transfers and store-job timeouts

@@ -579,7 +1041,8 @@ class ServerRole:

1041 # nothing else is in flight, finalize now so the peer

1042 # and the local store jobs don't wait for finish_request

1043 # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S.

582 - req = self._outbound.get(kv_request_id)

1044 + st = self._requests.get(kv_request_id)

1045 + req = st.outbound if st is not None else None

1046 if req is not None:

1047 req.finishing = True

1048 if not self._has_inflight_for(kv_request_id):

vllm/v1/kv_offload/tiering/p2p/session/session.py

+120 -19

@@ -34,6 +34,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import (

34 ConnectMsg,

35 DisconnectMsg,

36 FetchMsg,

37 + LookupMsg,

38 + LookupRespMsg,

39 TransferDoneMsg,

40 )

41 from vllm.v1.kv_offload.tiering.p2p.session.server import (

@@ -42,7 +44,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.server import (

44 )

46 if TYPE_CHECKING:

45 - from vllm.v1.kv_offload.tiering.base import JobId

47 + from vllm.v1.kv_offload.base import ReqContext

48 + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager

49 from vllm.v1.kv_offload.tiering.p2p.data import DataTransport

51 logger = init_logger(__name__)

@@ -71,6 +74,22 @@ class SessionPollResult(NamedTuple):

74 new_fetch_ids: list[str]

77 +class SessionCloseResult(NamedTuple):

78 + """Result of tearing down a P2PSession.

79 +

80 + `failed_jobs`/`failed_stores` are the in-flight jobs (client loads /

81 + server stores) the manager must mark failed. `failed_req_ids` is every

82 + client-side kv_request_id whose lookup() must fail (in-flight loads plus

83 + unresolved probes); `failed_serves` is the server-side lookup state the

84 + dead peer can no longer resolve.

85 + """

86 +

87 + failed_jobs: list[int] # client load job_ids

88 + failed_req_ids: list[str] # client kv_request_ids (loads + probes)

89 + failed_stores: list[int] # server store job_ids

90 + failed_serves: list[ReqContext] # server-side lookup ctxs needing release

91 +

92 +

93 class P2PSession:

94 """Bidirectional session — coordinator over ClientRole + ServerRole.

@@ -93,12 +112,14 @@ class P2PSession:

112 local_id: str,

113 transport: DataTransport,

114 local_block_len: int,

115 + local_hash_seed: str,

116 conn: ControlConnection | None = None,

117 ) -> None:

118 self.peer_id = peer_id

119 self._local_id = local_id

120 self._transport = transport

121 self._local_block_len = local_block_len

122 + self._local_hash_seed = local_hash_seed

123 self._conn: ControlConnection | None = None

124

125 self._send_ready = False # True after the peer acked our ConnectMsg

@@ -115,7 +136,11 @@ class P2PSession:

136 self._new_fetch_ids: list[str] = []

137

138 self._client = ClientRole(peer_id=peer_id, send=self._send)

118 - self._server = ServerRole(peer_id=peer_id, transport=transport, send=self._send)

139 + self._server = ServerRole(

140 + peer_id=peer_id,

141 + transport=transport,

142 + send=self._send,

143 + )

144

145 if conn is not None:

146 self.attach_connection(conn)

@@ -139,6 +164,11 @@ class P2PSession:

164 """True after the peer acked our ConnectMsg (we may send freely)."""

165 return self._send_ready

166

167 + @property

168 + def has_pending_work(self) -> bool:

169 + """True while inbound loads or outbound transfers are outstanding."""

170 + return self._client.has_active_loads or self._server.has_inflight_transfers

171 +

172 # ------------------------------------------------------------------

173 # Connection lifecycle

174 # ------------------------------------------------------------------

@@ -162,7 +192,7 @@ class P2PSession:

192 self,

193 job_id: JobId,

194 kv_request_id: str,

165 - keys: Sequence[bytes],

195 + keys: Sequence[OffloadKey],

196 block_ids: Sequence[int],

197 ) -> None:

198 """Send fetch to the peer."""

@@ -183,13 +213,40 @@ class P2PSession:

213 def finish_request(self, kv_request_id: str) -> None:

214 """Called when the request is finishing locally.

215

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.

216 + Finishes the client role (aborts any inbound load and drops any

217 + pending symmetric-P2P lookup state) and finalizes any outbound

218 + serving (server role) for this id. Roles that aren't active for

219 + this id are silent no-ops.

220 """

190 - self._client.cancel(kv_request_id)

221 + self._client.finish(kv_request_id)

222 self._server.finish(kv_request_id)

223

224 + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None:

225 + """Register or resolve one (kv_request_id, key) probe.

226 +

227 + Called from the manager's lookup() for symmetric-P2P consumers

228 + (``remote_kv_source`` sub-dict in kv_transfer_params). See

229 + ``ClientRole.register_lookup`` for the state-machine contract.

230 + """

231 + return self._client.register_lookup(kv_request_id, key)

232 +

233 + def flush_pending_lookups(self) -> None:

234 + """Flush any aggregated symmetric-P2P lookups for this peer.

235 +

236 + Called once per scheduler step from the manager's

237 + ``on_schedule_end()``. Send-gating is handled inside the

238 + client's ``_send`` callback (queues until ConnectAckMsg).

239 + """

240 + self._client.flush_pending_lookups()

241 +

242 + def serve_external_requests(self, parent: ParentManager) -> None:

243 + """Resolve inbound peer lookups against the tiering manager.

244 +

245 + Delegates to the server role; the ``parent`` handle is valid

246 + only for the duration of this call.

247 + """

248 + self._server.serve_external_requests(parent)

249 +

250 def poll(self) -> SessionPollResult:

251 """Process incoming messages, drive transfers, apply timeouts."""

252 if self._conn is None:

@@ -214,14 +271,22 @@ class P2PSession:

271 loads=loads, stores=stores, new_fetch_ids=new_fetch_ids

272 )

273

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.

274 + def close(self) -> SessionCloseResult:

275 + """Shut down.

276 +

277 + failed_jobs: client load job_ids to fail.

278 + failed_req_ids: client kv_request_ids to fail — in-flight loads plus

279 + requests with an unresolved symmetric-P2P probe toward the

280 + now-dead peer. The manager fails these so the consumer's lookup()

281 + falls back to local prefill instead of deferring forever on an

282 + answer that can never arrive.

283 + failed_stores: server store job_ids to fail.

284 + failed_serves: synthetic lookup ctxs still owing

285 + ``parent.on_request_finished`` (the manager flushes these on

286 + its next ``serve_external_requests``).

287 """

223 - failed_loads = self._client.close()

224 - failed_stores = self._server.close()

288 + client_result = self._client.close()

289 + failed_stores, failed_serves = self._server.close()

290

291 if self._conn is not None:

292 with contextlib.suppress(Exception):

@@ -229,7 +294,12 @@ class P2PSession:

294 self._conn.close()

295 self._conn = None

296

232 - return failed_loads, failed_stores

297 + return SessionCloseResult(

298 + failed_jobs=client_result.failed_jobs,

299 + failed_req_ids=client_result.failed_req_ids,

300 + failed_stores=failed_stores,

301 + failed_serves=failed_serves,

302 + )

303

304 # ------------------------------------------------------------------

305 # Message dispatch

@@ -299,9 +369,9 @@ class P2PSession:

369 elif msg_type == FetchMsg.TYPE:

370 FetchMsg.validate(msg)

371 kv_request_id = msg[FetchMsg.KV_REQUEST_ID]

302 - block_hashes = [

372 + keys = [

373 OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))

304 - for bh in msg[FetchMsg.BLOCK_HASHES]

374 + for bh in msg[FetchMsg.KEYS]

375 ]

376 block_indexes = msg[FetchMsg.BLOCK_INDEXES]

377 # Run the server-role state machine inline as today —

@@ -310,7 +380,7 @@ class P2PSession:

380 # the manager (after poll() returns) can replay any parked

381 # submit_store batches; their add_stored_blocks calls hit

382 # the demand recorded here and submit transfers immediately.

313 - self._server.on_fetch(kv_request_id, block_hashes, block_indexes)

383 + self._server.on_fetch(kv_request_id, keys, block_indexes)

384 self._new_fetch_ids.append(kv_request_id)

385 elif msg_type == AbortFetchMsg.TYPE:

386 AbortFetchMsg.validate(msg)

@@ -324,6 +394,23 @@ class P2PSession:

394 elif msg_type == AbortAckMsg.TYPE:

395 AbortAckMsg.validate(msg)

396 self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID])

397 + elif msg_type == LookupMsg.TYPE:

398 + LookupMsg.validate(msg)

399 + kv_request_id = msg[LookupMsg.KV_REQUEST_ID]

400 + keys = [

401 + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))

402 + for bh in msg[LookupMsg.KEYS]

403 + ]

404 + self._server.on_lookup(kv_request_id, keys)

405 + elif msg_type == LookupRespMsg.TYPE:

406 + LookupRespMsg.validate(msg)

407 + kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID]

408 + keys = [

409 + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))

410 + for bh in msg[LookupRespMsg.KEYS]

411 + ]

412 + hits = msg[LookupRespMsg.HITS]

413 + self._client.on_lookup_resp(kv_request_id, keys, hits)

414 elif msg_type == DisconnectMsg.TYPE:

415 if self._conn is not None:

416 self._conn.mark_dead()

@@ -361,6 +448,12 @@ class P2PSession:

448 f"config fingerprint mismatch from {self.peer_id}: "

449 f"remote={remote_fp!r}, local={local_fp!r}"

450 )

451 + if msg[ConnectMsg.HASH_SEED] != self._local_hash_seed:

452 + raise ValueError(

453 + f"PYTHONHASHSEED mismatch from {self.peer_id}: "

454 + f"remote={msg[ConnectMsg.HASH_SEED]!r}, "

455 + f"local={self._local_hash_seed!r}"

456 + )

457 self._transport.add_remote_peer(

458 self.peer_id,

459 agent_metadata=msg[ConnectMsg.AGENT_METADATA],

@@ -410,6 +503,7 @@ class P2PSession:

503 ConnectMsg.NUM_BLOCKS: self._transport.num_blocks,

504 ConnectMsg.BLOCK_LEN: self._transport.block_len,

505 ConnectMsg.CONFIG_FINGERPRINT: self._transport.config_fingerprint,

506 + ConnectMsg.HASH_SEED: self._local_hash_seed,

507 }

508 )

509

@@ -433,8 +527,15 @@ class P2PSession:

527 self._conn.send(msg)

528 logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY))

529 except Exception:

530 + # A send failure means the connection is broken. Swallowing it

531 + # silently strands every in-flight lookup/load toward this peer:

532 + # the session stays alive, is never reaped, and the consumer's

533 + # lookup() keeps returning RETRY until the HTTP client times out.

534 + # Mark the connection dead so the manager reaps the session on

535 + # its next poll and surfaces the stranded work as failures.

536 logger.warning(

437 - "P2PSession %s: failed to send %s",

537 + "P2PSession %s: send of %s failed — marking connection dead",

538 self.peer_id,

539 msg.get(TYPE_KEY),

540 )

541 + self._conn.mark_dead()