@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 30217b0e80

[Bugfix][KV Offload][P2P] Scope serve state to fetch rounds (#49877)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <Itay.etelis@gmail.com>

Itay Etelis committed Jul 28, 2026 at 18:01 UTC 30217b0e809ef045b3b4aa7c011d6c6d2b0b0362
6 files changed +573 -393
tests/v1/kv_offload/tiering/p2p/test_manager.py
+7 -6
@@ -1612,16 +1612,17 @@ class TestBindHostPortDefaults:
1612 monkeypatch.setattr(
1613 manager_module,
1614 "NixlTransport",
1615 - lambda agent_name, *a, **k: calls.update(nixl_name=agent_name)
1616 - or SimpleNamespace(),
1615 + lambda agent_name, *a, **k: (
1616 + calls.update(nixl_name=agent_name) or SimpleNamespace()
1617 + ),
1618 )
1619 monkeypatch.setattr(
1620 manager_module,
1621 "ZmqTransport",
1621 - lambda local_id, host, port, *a, **k: calls.update(
1622 - zmq_id=local_id, zmq_host=host, zmq_port=port
1623 - )
1624 - or SimpleNamespace(),
1622 + lambda local_id, host, port, *a, **k: (
1623 + calls.update(zmq_id=local_id, zmq_host=host, zmq_port=port)
1624 + or SimpleNamespace()
1625 + ),
1626 )
1627 spec = SimpleNamespace(
1628 blocks_per_chunk=1,
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+100 -16
@@ -49,6 +49,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
49 from vllm.v1.kv_offload.tiering.p2p.session.server import (
50 _CANCEL_DRAIN_TIMEOUT_S,
51 _InflightXfer,
52 + _OutboundRequestState,
53 )
54 from vllm.v1.kv_offload.tiering.p2p.session.session import (
55 _MAX_CONSECUTIVE_DISPATCH_ERRORS,
@@ -315,10 +316,26 @@ def _activate(
316 # either a missing entry or a None field. These helpers paper over that.
317
318
319 +def _client_load(session: P2PSession, kv_request_id: str):
320 + """The single in-flight load of a kv_request_id (loads are per-round)."""
321 + loads = session._client._requests[kv_request_id].loads
322 + assert len(loads) == 1
323 + return next(iter(loads.values()))
324 +
325 +
326 def _srv_outbound(session: P2PSession, kv_request_id: str):
319 - """Outbound serve state for a kv_request_id, or None (idle / GC'd)."""
327 + """Serve-side round for a kv_request_id, or None (idle / GC'd).
328 +
329 + Rounds are keyed by wire round_seq; surfaces the demanded round when
330 + a fetch has bound one, else any parked supply round.
331 + """
332 st = session._server._requests.get(kv_request_id)
321 - return st.outbound if st is not None else None
333 + if st is None or not st.outbound:
334 + return None
335 + for rnd in st.outbound.values():
336 + if rnd.demand_received:
337 + return rnd
338 + return next(iter(st.outbound.values()))
339
340
341 def _srv_lookups(session: P2PSession) -> list:
@@ -330,8 +347,10 @@ def _srv_lookups(session: P2PSession) -> list:
347
348 def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None:
349 """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
350 + for (kv, _), started in session._server._pending_aborts.items():
351 + if kv == kv_request_id:
352 + return started
353 + return None
354
355
356 def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int:
@@ -479,6 +498,7 @@ class TestClientFlows:
498 conn.enqueue(
499 {
500 TYPE_KEY: TransferDoneMsg.TYPE,
501 + TransferDoneMsg.ROUND_SEQ: 0,
502 TransferDoneMsg.KV_REQUEST_ID: "req-1",
503 TransferDoneMsg.SUCCESS: True,
504 }
@@ -495,6 +515,7 @@ class TestClientFlows:
515 conn.enqueue(
516 {
517 TYPE_KEY: TransferDoneMsg.TYPE,
518 + TransferDoneMsg.ROUND_SEQ: 0,
519 TransferDoneMsg.KV_REQUEST_ID: "req-1",
520 TransferDoneMsg.SUCCESS: False,
521 }
@@ -538,6 +559,7 @@ class TestClientFlows:
559 conn.enqueue(
560 {
561 TYPE_KEY: TransferDoneMsg.TYPE,
562 + TransferDoneMsg.ROUND_SEQ: 0,
563 TransferDoneMsg.KV_REQUEST_ID: "req-1",
564 TransferDoneMsg.SUCCESS: True,
565 }
@@ -552,7 +574,7 @@ class TestClientFlows:
574 session.request_blocks(
575 job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
576 )
555 - session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0
577 + _client_load(session, "req-1").submitted_at = time.monotonic() - 60.0
578 session.poll()
579 abort = conn._sent[-1]
580 assert abort[TYPE_KEY] == AbortFetchMsg.TYPE
@@ -569,7 +591,7 @@ class TestClientFlows:
591 job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0]
592 )
593 # 1) Trip the load timeout to send AbortFetch and stamp aborted_at.
572 - session._client._requests["req-7"].load.submitted_at = (
594 + _client_load(session, "req-7").submitted_at = (
595 time.monotonic() - _LOAD_TIMEOUT_S - 1.0
596 )
597 loads = session.poll().loads
@@ -579,11 +601,11 @@ class TestClientFlows:
601 and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7"
602 for m in conn._sent
603 )
582 - assert session._client._requests["req-7"].load.aborted_at is not None
604 + assert _client_load(session, "req-7").aborted_at is not None
605
606 # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever
607 # arrived from the peer.
586 - session._client._requests["req-7"].load.aborted_at = (
608 + _client_load(session, "req-7").aborted_at = (
609 time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
610 )
611 loads = session.poll().loads
@@ -599,17 +621,18 @@ class TestClientFlows:
621 session.request_blocks(
622 job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0]
623 )
602 - session._client._requests["req-8"].load.submitted_at = (
624 + _client_load(session, "req-8").submitted_at = (
625 time.monotonic() - _LOAD_TIMEOUT_S - 1.0
626 )
627 # First poll: AbortFetch goes out.
628 session.poll()
607 - assert session._client._requests["req-8"].load.aborted_at is not None
629 + assert _client_load(session, "req-8").aborted_at is not None
630
631 # Peer acks the abort.
632 conn.enqueue(
633 {
634 TYPE_KEY: AbortAckMsg.TYPE,
635 + AbortAckMsg.ROUND_SEQ: 0,
636 AbortAckMsg.KV_REQUEST_ID: "req-8",
637 }
638 )
@@ -915,6 +938,7 @@ class TestLookupFlow:
938 conn.enqueue(
939 {
940 TYPE_KEY: LookupMsg.TYPE,
941 + LookupMsg.ROUND_SEQ: 0,
942 LookupMsg.KV_REQUEST_ID: "req-1",
943 LookupMsg.KEYS: [b"hX", b"hY", b"hZ"],
944 }
@@ -951,6 +975,7 @@ def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]):
975 conn.enqueue(
976 {
977 TYPE_KEY: LookupMsg.TYPE,
978 + LookupMsg.ROUND_SEQ: 0,
979 LookupMsg.KV_REQUEST_ID: kv_request_id,
980 LookupMsg.KEYS: list(keys),
981 }
@@ -1201,6 +1226,7 @@ class TestServerLookupHandling:
1226 conn.enqueue(
1227 {
1228 TYPE_KEY: FetchMsg.TYPE,
1229 + FetchMsg.ROUND_SEQ: 0,
1230 FetchMsg.KV_REQUEST_ID: "req-1",
1231 FetchMsg.KEYS: [],
1232 FetchMsg.BLOCK_INDEXES: [],
@@ -1240,6 +1266,7 @@ class TestServerLookupHandling:
1266 conn.enqueue(
1267 {
1268 TYPE_KEY: FetchMsg.TYPE,
1269 + FetchMsg.ROUND_SEQ: 0,
1270 FetchMsg.KV_REQUEST_ID: "req-1",
1271 FetchMsg.KEYS: [b"hA", b"hB"],
1272 FetchMsg.BLOCK_INDEXES: [20, 21],
@@ -1277,6 +1304,7 @@ class TestServerFlows:
1304 conn.enqueue(
1305 {
1306 TYPE_KEY: FetchMsg.TYPE,
1307 + FetchMsg.ROUND_SEQ: 0,
1308 FetchMsg.KV_REQUEST_ID: "req-1",
1309 FetchMsg.KEYS: [b"k1", b"k2"],
1310 FetchMsg.BLOCK_INDEXES: [10, 11],
@@ -1295,6 +1323,7 @@ class TestServerFlows:
1323 conn.enqueue(
1324 {
1325 TYPE_KEY: FetchMsg.TYPE,
1326 + FetchMsg.ROUND_SEQ: 0,
1327 FetchMsg.KV_REQUEST_ID: "req-1",
1328 FetchMsg.KEYS: [b"k1"],
1329 FetchMsg.BLOCK_INDEXES: [5],
@@ -1313,6 +1342,7 @@ class TestServerFlows:
1342 conn.enqueue(
1343 {
1344 TYPE_KEY: FetchMsg.TYPE,
1345 + FetchMsg.ROUND_SEQ: 0,
1346 FetchMsg.KV_REQUEST_ID: "req-1",
1347 FetchMsg.KEYS: [b"k1"],
1348 FetchMsg.BLOCK_INDEXES: [5],
@@ -1331,6 +1361,7 @@ class TestServerFlows:
1361 conn.enqueue(
1362 {
1363 TYPE_KEY: AbortFetchMsg.TYPE,
1364 + AbortFetchMsg.ROUND_SEQ: 0,
1365 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1366 }
1367 )
@@ -1349,13 +1380,19 @@ class TestServerFlows:
1380 tid = 42
1381 session._server._inflight_add(
1382 tid,
1352 - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
1383 + _InflightXfer(
1384 + kv_request_id="req-1",
1385 + block_count=1,
1386 + job_ids={1},
1387 + round=_OutboundRequestState(inflight=1),
1388 + ),
1389 )
1390 transport._cancel_still_inflight.add(tid)
1391
1392 conn.enqueue(
1393 {
1394 TYPE_KEY: AbortFetchMsg.TYPE,
1395 + AbortFetchMsg.ROUND_SEQ: 0,
1396 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1397 }
1398 )
@@ -1376,13 +1413,19 @@ class TestServerFlows:
1413 tid = 42
1414 session._server._inflight_add(
1415 tid,
1379 - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
1416 + _InflightXfer(
1417 + kv_request_id="req-1",
1418 + block_count=1,
1419 + job_ids={1},
1420 + round=_OutboundRequestState(inflight=1),
1421 + ),
1422 )
1423 transport._cancel_still_inflight.add(tid)
1424
1425 conn.enqueue(
1426 {
1427 TYPE_KEY: AbortFetchMsg.TYPE,
1428 + AbortFetchMsg.ROUND_SEQ: 0,
1429 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1430 }
1431 )
@@ -1409,20 +1452,26 @@ class TestServerFlows:
1452 tid = 42
1453 session._server._inflight_add(
1454 tid,
1412 - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
1455 + _InflightXfer(
1456 + kv_request_id="req-1",
1457 + block_count=1,
1458 + job_ids={1},
1459 + round=_OutboundRequestState(inflight=1),
1460 + ),
1461 )
1462 transport._cancel_still_inflight.add(tid)
1463
1464 conn.enqueue(
1465 {
1466 TYPE_KEY: AbortFetchMsg.TYPE,
1467 + AbortFetchMsg.ROUND_SEQ: 0,
1468 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1469 }
1470 )
1471 session.poll()
1472 assert _srv_abort_started(session, "req-1") is not None
1473 # Backdate past the drain deadline.
1425 - session._server._requests["req-1"].abort_started_at = (
1474 + session._server._pending_aborts[("req-1", 0)] = (
1475 time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0
1476 )
1477 # Even if the transport still claims it can't cancel, the
@@ -1445,13 +1494,19 @@ class TestServerFlows:
1494 tid = 42
1495 session._server._inflight_add(
1496 tid,
1448 - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
1497 + _InflightXfer(
1498 + kv_request_id="req-1",
1499 + block_count=1,
1500 + job_ids={1},
1501 + round=_OutboundRequestState(inflight=1),
1502 + ),
1503 )
1504 transport._cancel_still_inflight.add(tid)
1505
1506 conn.enqueue(
1507 {
1508 TYPE_KEY: AbortFetchMsg.TYPE,
1509 + AbortFetchMsg.ROUND_SEQ: 0,
1510 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1511 }
1512 )
@@ -1463,6 +1518,7 @@ class TestServerFlows:
1518 conn.enqueue(
1519 {
1520 TYPE_KEY: AbortFetchMsg.TYPE,
1521 + AbortFetchMsg.ROUND_SEQ: 0,
1522 AbortFetchMsg.KV_REQUEST_ID: "req-1",
1523 }
1524 )
@@ -1497,6 +1553,7 @@ class TestServerFlows:
1553 conn.enqueue(
1554 {
1555 TYPE_KEY: FetchMsg.TYPE,
1556 + FetchMsg.ROUND_SEQ: 0,
1557 FetchMsg.KV_REQUEST_ID: "req-1",
1558 FetchMsg.KEYS: [b"k1"],
1559 FetchMsg.BLOCK_INDEXES: [5],
@@ -1529,6 +1586,7 @@ class TestServerFlows:
1586 conn.enqueue(
1587 {
1588 TYPE_KEY: FetchMsg.TYPE,
1589 + FetchMsg.ROUND_SEQ: 0,
1590 FetchMsg.KV_REQUEST_ID: "req-1",
1591 FetchMsg.KEYS: [b"k1"],
1592 FetchMsg.BLOCK_INDEXES: [5],
@@ -1571,6 +1629,7 @@ class TestFinishRequestServerSide:
1629 conn.enqueue(
1630 {
1631 TYPE_KEY: FetchMsg.TYPE,
1632 + FetchMsg.ROUND_SEQ: 0,
1633 FetchMsg.KV_REQUEST_ID: "req-1",
1634 FetchMsg.KEYS: [b"k1"],
1635 FetchMsg.BLOCK_INDEXES: [5],
@@ -1596,6 +1655,7 @@ class TestFinishRequestServerSide:
1655 conn.enqueue(
1656 {
1657 TYPE_KEY: FetchMsg.TYPE,
1658 + FetchMsg.ROUND_SEQ: 0,
1659 FetchMsg.KV_REQUEST_ID: "req-1",
1660 FetchMsg.KEYS: [b"k1", b"k2"],
1661 FetchMsg.BLOCK_INDEXES: [10, 11],
@@ -1631,6 +1691,7 @@ class TestFinishRequestServerSide:
1691 conn.enqueue(
1692 {
1693 TYPE_KEY: FetchMsg.TYPE,
1694 + FetchMsg.ROUND_SEQ: 0,
1695 FetchMsg.KV_REQUEST_ID: "req-1",
1696 FetchMsg.KEYS: [b"k1"],
1697 FetchMsg.BLOCK_INDEXES: [10],
@@ -1665,6 +1726,7 @@ class TestFinishRequestServerSide:
1726 conn.enqueue(
1727 {
1728 TYPE_KEY: FetchMsg.TYPE,
1729 + FetchMsg.ROUND_SEQ: 0,
1730 FetchMsg.KV_REQUEST_ID: "req-1",
1731 FetchMsg.KEYS: [b"k1"],
1732 FetchMsg.BLOCK_INDEXES: [10],
@@ -1686,6 +1748,7 @@ class TestFinishRequestServerSide:
1748 conn.enqueue(
1749 {
1750 TYPE_KEY: FetchMsg.TYPE,
1751 + FetchMsg.ROUND_SEQ: 0,
1752 FetchMsg.KV_REQUEST_ID: "req-2",
1753 FetchMsg.KEYS: [b"k1", b"k2"],
1754 FetchMsg.BLOCK_INDEXES: [10, 11],
@@ -1719,6 +1782,7 @@ class TestFinishRequestServerSide:
1782 conn.enqueue(
1783 {
1784 TYPE_KEY: FetchMsg.TYPE,
1785 + FetchMsg.ROUND_SEQ: 0,
1786 FetchMsg.KV_REQUEST_ID: "req-1",
1787 FetchMsg.KEYS: [b"demand"],
1788 FetchMsg.BLOCK_INDEXES: [5],
@@ -1753,6 +1817,7 @@ class TestFinishRequestServerSide:
1817 conn.enqueue(
1818 {
1819 TYPE_KEY: FetchMsg.TYPE,
1820 + FetchMsg.ROUND_SEQ: 0,
1821 FetchMsg.KV_REQUEST_ID: "req-1",
1822 FetchMsg.KEYS: [b"k1"],
1823 FetchMsg.BLOCK_INDEXES: [10],
@@ -1790,6 +1855,7 @@ class TestFinishRequestServerSide:
1855 conn.enqueue(
1856 {
1857 TYPE_KEY: FetchMsg.TYPE,
1858 + FetchMsg.ROUND_SEQ: 0,
1859 FetchMsg.KV_REQUEST_ID: "req-1",
1860 FetchMsg.KEYS: [b"k1"],
1861 FetchMsg.BLOCK_INDEXES: [10],
@@ -1824,6 +1890,7 @@ class TestFinishRequestServerSide:
1890 conn.enqueue(
1891 {
1892 TYPE_KEY: FetchMsg.TYPE,
1893 + FetchMsg.ROUND_SEQ: 0,
1894 FetchMsg.KV_REQUEST_ID: "req-1",
1895 FetchMsg.KEYS: [b"k1", b"k2", b"k3"],
1896 FetchMsg.BLOCK_INDEXES: [10, 11, 12],
@@ -1885,6 +1952,7 @@ class TestFinishRequestServerSide:
1952 conn.enqueue(
1953 {
1954 TYPE_KEY: FetchMsg.TYPE,
1955 + FetchMsg.ROUND_SEQ: 0,
1956 FetchMsg.KV_REQUEST_ID: "req-1",
1957 FetchMsg.KEYS: [b"k1", b"k2"],
1958 FetchMsg.BLOCK_INDEXES: [10, 11],
@@ -1951,6 +2019,7 @@ class TestBidirectional:
2019 conn.enqueue(
2020 {
2021 TYPE_KEY: FetchMsg.TYPE,
2022 + FetchMsg.ROUND_SEQ: 0,
2023 FetchMsg.KV_REQUEST_ID: "req-srv",
2024 FetchMsg.KEYS: [b"served"],
2025 FetchMsg.BLOCK_INDEXES: [7],
@@ -1981,6 +2050,7 @@ class TestBidirectional:
2050 conn.enqueue(
2051 {
2052 TYPE_KEY: TransferDoneMsg.TYPE,
2053 + TransferDoneMsg.ROUND_SEQ: 0,
2054 TransferDoneMsg.KV_REQUEST_ID: "req-cli",
2055 TransferDoneMsg.SUCCESS: True,
2056 }
@@ -2168,6 +2238,7 @@ class TestAdversarial:
2238 conn.enqueue(
2239 {
2240 TYPE_KEY: FetchMsg.TYPE,
2241 + FetchMsg.ROUND_SEQ: 0,
2242 FetchMsg.KV_REQUEST_ID: "req-bad",
2243 FetchMsg.KEYS: [b"k1", b"k2"],
2244 FetchMsg.BLOCK_INDEXES: [1],
@@ -2216,6 +2287,7 @@ class TestDispatchErrorHandling:
2287 conn.enqueue(
2288 {
2289 TYPE_KEY: FetchMsg.TYPE,
2290 + FetchMsg.ROUND_SEQ: 0,
2291 FetchMsg.KV_REQUEST_ID: "req-bad",
2292 FetchMsg.KEYS: [b"k1", b"k2"],
2293 FetchMsg.BLOCK_INDEXES: [1],
@@ -2246,6 +2318,7 @@ class TestDispatchErrorHandling:
2318 conn.enqueue(
2319 {
2320 TYPE_KEY: FetchMsg.TYPE,
2321 + FetchMsg.ROUND_SEQ: 0,
2322 FetchMsg.KV_REQUEST_ID: "req-1",
2323 FetchMsg.KEYS: [b"k1"],
2324 FetchMsg.BLOCK_INDEXES: [0],
@@ -2269,6 +2342,7 @@ class TestDispatchErrorHandling:
2342 conn.enqueue(
2343 {
2344 TYPE_KEY: FetchMsg.TYPE,
2345 + FetchMsg.ROUND_SEQ: 0,
2346 FetchMsg.KV_REQUEST_ID: "req-1",
2347 FetchMsg.KEYS: [b"k1"],
2348 FetchMsg.BLOCK_INDEXES: [0],
@@ -2296,6 +2370,7 @@ class TestDispatchErrorHandling:
2370 conn.enqueue(
2371 {
2372 TYPE_KEY: FetchMsg.TYPE,
2373 + FetchMsg.ROUND_SEQ: 0,
2374 FetchMsg.KV_REQUEST_ID: "req-1",
2375 FetchMsg.KEYS: [b"k1"],
2376 FetchMsg.BLOCK_INDEXES: [0],
@@ -2341,6 +2416,7 @@ class TestInflightPerReqInvariant:
2416 conn.enqueue(
2417 {
2418 TYPE_KEY: FetchMsg.TYPE,
2419 + FetchMsg.ROUND_SEQ: 0,
2420 FetchMsg.KV_REQUEST_ID: kv_id,
2421 FetchMsg.KEYS: keys,
2422 FetchMsg.BLOCK_INDEXES: indexes,
@@ -2394,7 +2470,12 @@ class TestInflightPerReqInvariant:
2470 tid = kv_id_idx * 10 + j
2471 session._server._inflight_add(
2472 tid,
2397 - _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}),
2473 + _InflightXfer(
2474 + kv_request_id=kv_id,
2475 + block_count=1,
2476 + job_ids={tid},
2477 + round=_OutboundRequestState(inflight=1),
2478 + ),
2479 )
2480 assert _srv_total_inflight(session) == len(session._server._inflight)
2481 assert session._server._has_inflight_for("req-0")
@@ -2488,6 +2569,7 @@ class TestFetchMsgValidation:
2569 def _valid_msg(self) -> dict:
2570 return {
2571 TYPE_KEY: FetchMsg.TYPE,
2572 + FetchMsg.ROUND_SEQ: 0,
2573 FetchMsg.KV_REQUEST_ID: "req-1",
2574 FetchMsg.KEYS: [b"k1", b"k2"],
2575 FetchMsg.BLOCK_INDEXES: [0, 1],
@@ -2513,6 +2595,7 @@ class TestTransferDoneMsgValidation:
2595 def test_valid_message_passes(self):
2596 msg = {
2597 TYPE_KEY: TransferDoneMsg.TYPE,
2598 + TransferDoneMsg.ROUND_SEQ: 0,
2599 TransferDoneMsg.KV_REQUEST_ID: "req-1",
2600 TransferDoneMsg.SUCCESS: True,
2601 }
@@ -2521,6 +2604,7 @@ class TestTransferDoneMsgValidation:
2604 def test_success_wrong_type(self):
2605 msg = {
2606 TYPE_KEY: TransferDoneMsg.TYPE,
2607 + TransferDoneMsg.ROUND_SEQ: 0,
2608 TransferDoneMsg.KV_REQUEST_ID: "req-1",
2609 TransferDoneMsg.SUCCESS: 1,
2610 }
vllm/v1/kv_offload/tiering/p2p/session/client.py
+151 -146
@@ -11,7 +11,6 @@ callback injected by the coordinator (which gates on ConnectAck).
11
12 from __future__ import annotations
13
14 -import enum
14 import time
15 from collections.abc import Callable, Sequence
16 from dataclasses import dataclass, field
@@ -35,26 +34,12 @@ _LOAD_TIMEOUT_S = 30.0
34 _ABORT_ACK_TIMEOUT_S = 10.0
35
36
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 -
37 @dataclass
38 class _InboundLoadState:
39 """Client-role state for a single in-flight load request.
40
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.
41 + Lives in ``_ClientRequestState.loads`` keyed by round_seq for the
42 + duration of a fetch; the owning kv_request_id is the outer dict key.
43 """
44
45 job_id: int # opaque ID assigned by the manager to this load request
@@ -68,7 +53,7 @@ class _ClientRequestState:
53
54 One entry per kv_request_id we're driving. Lookup-phase fields are
55 used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave
71 - ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An
56 + ``probes``/``unsent`` empty and drive just ``phase`` and ``loads``. An
57 entry is dropped once every field is idle — see ``ClientRole._maybe_prune``.
58 """
59
@@ -82,11 +67,23 @@ class _ClientRequestState:
67 # OffloadKeys registered but not yet flushed onto the wire. Drained and
68 # cleared by the next flush_pending_lookups.
69 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
70 + # Current lookup round. LookupMsgs carry it, each fetch closes it and
71 + # advances it, so every round's supply/demand/completion is isolated
72 + # on the wire. PD clients never probe and stay on round 0.
73 + round_seq: int = 0
74 + # This id ran the symmetric lookup phase (register_lookup); a fetch
75 + # with keys then requires every key to be a confirmed probe. PD
76 + # loads never probe.
77 + probed: bool = False
78 +
79 + # The peer holds lookup state no FetchMsg has closed: a LookupMsg
80 + # was flushed since the last fetch. finish owes a terminal empty
81 + # FetchMsg while set, so the peer releases parked supply.
82 + peer_lookup_open: bool = False
83 + # In-flight loads keyed by the round their fetch carried. The
84 + # scheduler submits loads incrementally as chunks resolve, so several
85 + # can be in flight at once; TransferDone/AbortAck match by round.
86 + loads: dict[int, _InboundLoadState] = field(default_factory=dict)
87
88
89 class LoadResult(NamedTuple):
@@ -133,11 +130,10 @@ class ClientRole:
130 # _serve_pending. Populated by register_lookup, drained by
131 # flush_pending_lookups, and discarded on finish/close.
132 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.
133 + # kv_request_ids with at least one fetch in flight — the work-list
134 + # collect_results walks for timeouts, and the has_active_loads
135 + # predicate, instead of scanning every request. Kept in exact sync
136 + # with ``st.loads``.
137 self._active_loads: set[str] = set()
138 self._completed_loads: list[LoadResult] = []
139
@@ -156,17 +152,21 @@ class ClientRole:
152 def _maybe_prune(self, kv_request_id: str) -> None:
153 """Drop the entry once it holds no live load or lookup state.
154
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.
155 + ``peer_lookup_open`` is only read by ``finish``, and every path
156 + that clears the last probe (fetch / finish / close) also settles
157 + it, so dropping on emptiness never loses a flag still in use.
158 """
159 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:
160 + if st is not None and not st.loads and not st.probes and not st.unsent:
161 del self._requests[kv_request_id]
162
163 + def _on_load_terminal(self, kv_request_id: str, st: _ClientRequestState) -> None:
164 + """Wind down id-level state once no load remains in flight."""
165 + if st.loads:
166 + return
167 + self._active_loads.discard(kv_request_id)
168 + self._maybe_prune(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."""
@@ -184,7 +184,12 @@ class ClientRole:
184 block_ids: Sequence[int],
185 send_ready: bool,
186 ) -> None:
187 - """Register a load request and send the FetchMsg."""
187 + """Send the FetchMsg closing the current lookup round.
188 +
189 + The scheduler may submit several loads per kv_request_id as its
190 + matched prefix resolves incrementally; each fetch carries the
191 + round it closes so the loads stay independent on the wire.
192 + """
193 logger.debug(
194 "P2PSession %s: request_blocks job_id=%d kv_request_id=%s "
195 "blocks=%d ready=%s",
@@ -195,77 +200,74 @@ class ClientRole:
200 send_ready,
201 )
202 st = self._get_or_create_request(kv_request_id)
198 - st.load = _InboundLoadState(
203 + round_seq = st.round_seq
204 + st.round_seq += 1
205 + st.loads[round_seq] = _InboundLoadState(
206 job_id=job_id,
207 submitted_at=time.monotonic(),
208 )
209 self._active_loads.add(kv_request_id)
203 - st.phase = ClientPhase.FETCH_SENT
210 + st.peer_lookup_open = False
211 self._send(
212 {
213 TYPE_KEY: FetchMsg.TYPE,
214 FetchMsg.KV_REQUEST_ID: kv_request_id,
215 FetchMsg.KEYS: list(keys),
216 FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids],
217 + FetchMsg.ROUND_SEQ: round_seq,
218 }
219 )
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:
220 + # Issuing the fetch closes this lookup round, so drop all probe
221 + # state. Once the peer serves the fetch both sides unpin, so a
222 + # stale cached True would let a later lookup() return HIT for a
223 + # block the producer may have evicted; clearing forces a fresh
224 + # probe under the next round.
225 + if st.probed and keys:
226 + assert st.probes, (
227 + f"symmetric fetch for {kv_request_id} has keys but no probes"
228 + )
229 assert all(st.probes.get(key) is True for key in keys)
230 st.probes.clear()
231
232 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 -
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.
233 + """Finish a request: abort in-flight loads and release lookup state.
234 +
235 + Called from the session's ``finish_request``. Sends an
236 + AbortFetchMsg per load not already aborting, and — independently —
237 + the terminal empty FetchMsg when the peer still holds lookup
238 + state no fetch has closed (its "request finished" signal: it
239 + releases lookup state, drains parked supply, and fires
240 + ``cb.finish_request``). A later round's supply can be parked
241 + while an earlier round's load is still in flight, so both can be
242 + owed at once.
243
244 Then drop all probe/lookup state and prune the entry.
245 """
246 st = self._requests.get(kv_request_id)
247 if st is None:
248 return
251 - if st.load is not None:
252 - if st.load.aborted_at is None:
249 + if st.loads:
250 + for round_seq, load in st.loads.items():
251 + if load.aborted_at is not None:
252 + continue
253 self._send(
254 {
255 TYPE_KEY: AbortFetchMsg.TYPE,
256 AbortFetchMsg.KV_REQUEST_ID: kv_request_id,
257 + AbortFetchMsg.ROUND_SEQ: round_seq,
258 }
259 )
259 - st.load = None
260 + st.loads.clear()
261 self._active_loads.discard(kv_request_id)
261 - elif st.phase is ClientPhase.PROBING:
262 - st.phase = ClientPhase.FETCH_SENT
262 + if st.peer_lookup_open:
263 + st.peer_lookup_open = False
264 self._send(
265 {
266 TYPE_KEY: FetchMsg.TYPE,
267 FetchMsg.KV_REQUEST_ID: kv_request_id,
268 FetchMsg.KEYS: [],
269 FetchMsg.BLOCK_INDEXES: [],
270 + FetchMsg.ROUND_SEQ: st.round_seq,
271 }
272 )
273 st.probes.clear()
@@ -273,20 +275,21 @@ class ClientRole:
275 self._flush_pending.discard(kv_request_id)
276 self._maybe_prune(kv_request_id)
277
276 - def on_transfer_done(self, kv_request_id: str, success: bool) -> None:
278 + def on_transfer_done(
279 + self, kv_request_id: str, success: bool, round_seq: int
280 + ) -> None:
281 """Handle a TransferDoneMsg from the peer."""
282 st = self._requests.get(kv_request_id)
279 - if st is not None and st.load is not None:
283 + load = st.loads.pop(round_seq, None) if st is not None else None
284 + if st is not None and load is not None:
285 self._completed_loads.append(
286 LoadResult(
282 - job_id=st.load.job_id,
287 + job_id=load.job_id,
288 kv_request_id=kv_request_id,
289 success=success,
290 )
291 )
287 - st.load = None
288 - self._active_loads.discard(kv_request_id)
289 - self._maybe_prune(kv_request_id)
292 + self._on_load_terminal(kv_request_id, st)
293 else:
294 # No matching in-flight load: either a duplicate
295 # transfer_done from the peer (protocol violation) or a
@@ -295,41 +298,44 @@ class ClientRole:
298 # so we can't tell — log so it's findable.
299 logger.warning(
300 "P2PSession %s: transfer_done for unknown kv_request_id=%s "
298 - "(duplicate from peer, or raced with local cancel/timeout)",
301 + "round=%s (duplicate from peer, or raced with local "
302 + "cancel/timeout)",
303 self._peer_id,
304 kv_request_id,
305 + round_seq,
306 )
307
303 - def on_abort_ack(self, kv_request_id: str) -> None:
308 + def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None:
309 """Handle an AbortAckMsg from the peer."""
310 st = self._requests.get(kv_request_id)
306 - if st is not None and st.load is not None:
311 + load = st.loads.pop(round_seq, None) if st is not None else None
312 + if st is not None and load is not None:
313 logger.warning(
314 "P2PSession %s: load request %s (job_id=%d) timed out; "
315 "load job completed with failure. If this recurs, ensure "
316 "PYTHONHASHSEED is set to the same value on all nodes.",
317 self._peer_id,
318 kv_request_id,
313 - st.load.job_id,
319 + load.job_id,
320 )
321 self._completed_loads.append(
322 LoadResult(
317 - job_id=st.load.job_id,
323 + job_id=load.job_id,
324 kv_request_id=kv_request_id,
325 success=False,
326 )
327 )
322 - st.load = None
323 - self._active_loads.discard(kv_request_id)
324 - self._maybe_prune(kv_request_id)
328 + self._on_load_terminal(kv_request_id, st)
329 else:
330 # See on_transfer_done: same ambiguity (duplicate ack
331 # vs. raced with local cancel/timeout that already popped).
332 logger.warning(
333 "P2PSession %s: abort_ack for unknown kv_request_id=%s "
330 - "(duplicate from peer, or raced with local cancel/timeout)",
334 + "round=%s (duplicate from peer, or raced with local "
335 + "cancel/timeout)",
336 self._peer_id,
337 kv_request_id,
338 + round_seq,
339 )
340
341 # ------------------------------------------------------------------
@@ -345,17 +351,18 @@ class ClientRole:
351 - Once a LookupRespMsg has resolved the entry: returns the cached
352 bool result on every call without popping it.
353
348 - A resolved entry is retained until its fetch is issued
349 - (``request_blocks`` pops it) or the request finishes
354 + A resolved entry is retained until a fetch closes the round
355 + (``request_blocks`` clears all probes) or the request finishes
356 (``finish`` clears all entries for the id). A request's
357 block set can be re-probed across steps, so popping on read would
358 make a repeat probe of an already-resolved key look brand-new and
359 re-queue it, emitting a redundant LookupMsg for an answer we
360 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.
361 + free; clearing at fetch forces a fresh probe under the next
362 + round, since the block is unpinned once served.
363 """
364 st = self._get_or_create_request(kv_request_id)
365 + st.probed = True
366 okey = OffloadKey(key)
367 if okey in st.probes:
368 return st.probes[okey]
@@ -378,14 +385,11 @@ class ClientRole:
385 ``on_schedule_end()``. A request's block set may be discovered
386 across several scheduler steps, so more than one LookupMsg can
387 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).
388 + keys, all tagged with the current round. register_lookup()
389 + de-dups in-flight and already-resolved (req_id, key) pairs, so
390 + each LookupMsg carries only the keys first probed in that step.
391 + Send-gating is handled by the injected ``_send`` callback
392 + (queues until ConnectAckMsg if needed).
393
394 Only requests that registered new keys since the last flush are
395 visited — the ``_flush_pending`` work-list avoids scanning every
@@ -395,13 +399,9 @@ class ClientRole:
399 st = self._requests.get(req_id)
400 if st is None or not st.unsent:
401 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
402 + # The peer now holds lookup state for this id; finish owes a
403 + # terminal empty FetchMsg until a fetch closes it.
404 + st.peer_lookup_open = True
405 logger.debug(
406 "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d",
407 self._peer_id,
@@ -413,6 +413,7 @@ class ClientRole:
413 TYPE_KEY: LookupMsg.TYPE,
414 LookupMsg.KV_REQUEST_ID: req_id,
415 LookupMsg.KEYS: list(st.unsent),
416 + LookupMsg.ROUND_SEQ: st.round_seq,
417 }
418 )
419 st.unsent = []
@@ -451,52 +452,56 @@ class ClientRole:
452 def collect_results(self) -> list[LoadResult]:
453 """Walk load timeouts and drain completed loads.
454
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.
455 + Loads past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and
456 + enter the aborting phase. Aborting loads past
457 + ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed.
458
459 Lookups have no timeout: an unanswered probe stays None (RETRY)
460 until finish_request clears it — see ``_ClientRequestState.probes``.
461 """
462 now = time.monotonic()
462 - to_remove: list[str] = []
463 + to_remove: list[tuple[str, int]] = []
464 for req_id in self._active_loads:
465 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,
473 - req_id,
474 - )
475 - self._send(
476 - {
477 - TYPE_KEY: AbortFetchMsg.TYPE,
478 - AbortFetchMsg.KV_REQUEST_ID: req_id,
479 - }
480 - )
481 - else:
482 - if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S:
483 - to_remove.append(req_id)
484 - self._completed_loads.append(
485 - LoadResult(
486 - job_id=load.job_id,
487 - kv_request_id=req_id,
488 - success=False,
466 + assert st.loads
467 + for round_seq, load in st.loads.items():
468 + if load.aborted_at is None:
469 + if now - load.submitted_at >= _LOAD_TIMEOUT_S:
470 + load.aborted_at = now
471 + logger.warning(
472 + "P2PSession %s: %s round=%s timed out, sending abort",
473 + self._peer_id,
474 + req_id,
475 + round_seq,
476 + )
477 + self._send(
478 + {
479 + TYPE_KEY: AbortFetchMsg.TYPE,
480 + AbortFetchMsg.KV_REQUEST_ID: req_id,
481 + AbortFetchMsg.ROUND_SEQ: round_seq,
482 + }
483 + )
484 + else:
485 + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S:
486 + to_remove.append((req_id, round_seq))
487 + self._completed_loads.append(
488 + LoadResult(
489 + job_id=load.job_id,
490 + kv_request_id=req_id,
491 + success=False,
492 + )
493 )
490 - )
491 - logger.warning(
492 - "P2PSession %s: abort_ack timed out for kv_request_id=%s",
493 - self._peer_id,
494 - req_id,
495 - )
496 - for req_id in to_remove:
497 - self._requests[req_id].load = None
498 - self._active_loads.discard(req_id)
499 - self._maybe_prune(req_id)
494 + logger.warning(
495 + "P2PSession %s: abort_ack timed out for "
496 + "kv_request_id=%s round=%s",
497 + self._peer_id,
498 + req_id,
499 + round_seq,
500 + )
501 + for req_id, round_seq in to_remove:
502 + st = self._requests[req_id]
503 + st.loads.pop(round_seq)
504 + self._on_load_terminal(req_id, st)
505
506 results = self._completed_loads
507 self._completed_loads = []
@@ -510,12 +515,12 @@ class ClientRole:
515 forever on an answer that can never arrive). See ``ClientCloseResult``.
516 """
517 failed_jobs = [
513 - st.load.job_id for st in self._requests.values() if st.load is not None
518 + load.job_id for st in self._requests.values() for load in st.loads.values()
519 ]
520 failed_req_ids = [
521 req_id
522 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())
523 + if st.loads or any(hit is None for hit in st.probes.values())
524 ]
525 self._requests.clear()
526 self._flush_pending.clear()
vllm/v1/kv_offload/tiering/p2p/session/protocol.py
+31 -19
@@ -26,14 +26,12 @@ Block Transfer Flow (happy path)
26
27 1. Client sends FetchMsg with a kv_request_id and lists of
28 block keys + remote indexes where it wants the data written.
29 - 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.
29 + A request may run several lookup→fetch rounds; symmetric-P2P
30 + messages carry ROUND_SEQ so each round's supply, demand, and
31 + completion stay isolated. The terminal empty FetchMsg is the
32 + server-side "request finished" signal for the id: parked
33 + LookupMsg batches are popped and ``cb.finish_request`` fires on
34 + each.
35 2. Server matches requested blocks against locally stored blocks:
36 - Blocks already available are transferred immediately via RDMA.
37 - Blocks not yet available are recorded as "demanded" and
@@ -178,33 +176,32 @@ class DisconnectMsg:
176
177
178 class FetchMsg:
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.
179 + """Client → Server: request blocks for one lookup round.
180 +
181 + A non-empty fetch closes only its round. The terminal empty FetchMsg
182 + (``KEYS`` and ``BLOCK_INDEXES`` both empty) is the "request
183 + finished" signal: the server pops parked LookupMsg batches, calls
184 + ``cb.finish_request`` on each, and drains any leftover supply.
185
186 Fields:
187 KV_REQUEST_ID: Identifies this block transfer request.
188 KEYS: List of block keys (OffloadKey bytes). May be empty.
189 BLOCK_INDEXES: List of remote block indexes (same length as KEYS).
190 + ROUND_SEQ: Lookup round this fetch closes. PD clients never probe
191 + and stay on their single round 0.
192 """
193
194 TYPE = "fetch"
195 KV_REQUEST_ID = "kv_request_id"
196 KEYS = "keys"
197 BLOCK_INDEXES = "block_indexes"
198 + ROUND_SEQ = "round_seq"
199
200 @staticmethod
201 def validate(msg: dict) -> None:
202 """Raise ValueError if any field has an invalid type or value."""
203 _require(msg, FetchMsg.KV_REQUEST_ID, str)
204 + _require_non_neg_int(msg, FetchMsg.ROUND_SEQ)
205 _require_list(msg, FetchMsg.KEYS)
206 _require_list(msg, FetchMsg.BLOCK_INDEXES)
207 keys = msg[FetchMsg.KEYS]
@@ -229,17 +226,21 @@ class LookupMsg:
226 Fields:
227 KV_REQUEST_ID: Identifies this lookup transaction.
228 KEYS: List of block keys (OffloadKey bytes) to probe.
229 + ROUND_SEQ: Lookup round these probes belong to; pinned supply is
230 + parked under it for that round's fetch.
231 """
232
233 TYPE = "lookup"
234 KV_REQUEST_ID = "kv_request_id"
235 KEYS = "keys"
236 + ROUND_SEQ = "round_seq"
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 + _require_non_neg_int(msg, LookupMsg.ROUND_SEQ)
244
245
246 class LookupRespMsg:
@@ -284,17 +285,22 @@ class TransferDoneMsg:
285 Fields:
286 KV_REQUEST_ID: The request that completed.
287 SUCCESS: Whether the transfer completed successfully.
288 + ROUND_SEQ: The fetch round that completed. Several loads can be
289 + in flight per id (the scheduler submits loads incrementally),
290 + so completions are matched by round.
291 """
292
293 TYPE = "transfer_done"
294 KV_REQUEST_ID = "kv_request_id"
295 SUCCESS = "success"
296 + ROUND_SEQ = "round_seq"
297
298 @staticmethod
299 def validate(msg: dict) -> None:
300 """Raise ValueError if any field has an invalid type or value."""
301 _require(msg, TransferDoneMsg.KV_REQUEST_ID, str)
302 _require(msg, TransferDoneMsg.SUCCESS, bool)
303 + _require_non_neg_int(msg, TransferDoneMsg.ROUND_SEQ)
304
305
306 class AbortFetchMsg:
@@ -302,15 +308,18 @@ class AbortFetchMsg:
308
309 Fields:
310 KV_REQUEST_ID: The request to cancel.
311 + ROUND_SEQ: The fetch round to cancel.
312 """
313
314 TYPE = "abort_fetch"
315 KV_REQUEST_ID = "kv_request_id"
316 + ROUND_SEQ = "round_seq"
317
318 @staticmethod
319 def validate(msg: dict) -> None:
320 """Raise ValueError if any field has an invalid type or value."""
321 _require(msg, AbortFetchMsg.KV_REQUEST_ID, str)
322 + _require_non_neg_int(msg, AbortFetchMsg.ROUND_SEQ)
323
324
325 class AbortAckMsg:
@@ -318,12 +327,15 @@ class AbortAckMsg:
327
328 Fields:
329 KV_REQUEST_ID: The request that was cancelled.
330 + ROUND_SEQ: The round that was cancelled; echoes AbortFetchMsg.
331 """
332
333 TYPE = "abort_ack"
334 KV_REQUEST_ID = "kv_request_id"
335 + ROUND_SEQ = "round_seq"
336
337 @staticmethod
338 def validate(msg: dict) -> None:
339 """Raise ValueError if any field has an invalid type or value."""
340 _require(msg, AbortAckMsg.KV_REQUEST_ID, str)
341 + _require_non_neg_int(msg, AbortAckMsg.ROUND_SEQ)
vllm/v1/kv_offload/tiering/p2p/session/server.py
+272 -202
@@ -53,15 +53,6 @@ class StoreResult(NamedTuple):
53 success: bool
54
55
56 -class _InflightXfer(NamedTuple):
57 - """Metadata for a single inflight RDMA transfer, keyed by transfer_id."""
58 -
59 - kv_request_id: str
60 - block_count: int
61 - # The set of store job IDs that contributed blocks to this transfer.
62 - job_ids: set[int]
63 -
64 -
56 class _MatchResult(NamedTuple):
57 """Result of block matching: pairs ready for transfer."""
58
@@ -73,12 +64,17 @@ class _MatchResult(NamedTuple):
64
65 @dataclass
66 class _OutboundRequestState:
76 - """Server-role state for a single peer fetch request.
67 + """Server-role state for a single fetch round of a peer request.
68
78 - The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key
79 - and is not duplicated on the value.
69 + A kv_request_id may run several lookup→fetch rounds; rounds live in
70 + ``_ServerRequestState.outbound`` keyed by the wire ``round_seq``, so
71 + terminals touch only their own round.
72 """
73
74 + # Supply came from inbound lookup pins (symmetric): no late
75 + # submit_store can arrive, so unmatched fetch demand fails fast. PD
76 + # rounds park demand for stores instead.
77 + lookup_supplied: bool = False
78 demand_received: bool = False
79 available: dict[OffloadKey, tuple[int, int]] = field(
80 default_factory=dict
@@ -88,7 +84,8 @@ class _OutboundRequestState:
84 ) # key → remote_block_idx: blocks peer wants, awaiting supply
85 remaining: int = 0 # blocks that need to be transferred to client
86 finishing: bool = False # Signal finish request ASAP
91 - # Job IDs that submit_store'd blocks for this request and have not
87 + inflight: int = 0 # transfers submitted for this round, not yet polled
88 + # Job IDs that submit_store'd blocks for this round and have not
89 # yet emitted a StoreResult. The terminal-finalize helper drains
90 # this set; poll-done and poll-failed discard entries as their
91 # StoreResults fire.
@@ -145,6 +142,21 @@ class _OutboundRequestState:
142 )
143
144
145 +@dataclass
146 +class _InflightXfer:
147 + """Metadata for a single inflight RDMA transfer, keyed by transfer_id."""
148 +
149 + kv_request_id: str
150 + block_count: int
151 + # The set of store job IDs that contributed blocks to this transfer.
152 + job_ids: set[int]
153 + # Round this transfer serves and its key in ``st.outbound``;
154 + # remaining/finalize apply only while the round is still registered.
155 + # Dummy default for test-seeded entries.
156 + round: _OutboundRequestState = field(default_factory=_OutboundRequestState)
157 + round_key: int = 0
158 +
159 +
160 @dataclass
161 class _ActiveLookup:
162 """In-flight state for one inbound LookupMsg.
@@ -159,6 +171,8 @@ class _ActiveLookup:
171 lookup_id: int
172 kv_request_id: str
173 ctx: ReqContext
174 + # Wire round these probes belong to; pins park under it.
175 + round_seq: int = 0
176 # Keys from the inbound LookupMsg, preserved in wire order so
177 # the aggregated response goes back in the same order.
178 keys: list[OffloadKey] = field(default_factory=list)
@@ -185,6 +199,7 @@ class _PendingLookup(NamedTuple):
199
200 keys: list[OffloadKey]
201 enqueued_at: float
202 + round_seq: int = 0
203
204
205 @dataclass
@@ -198,17 +213,15 @@ class _ServerRequestState:
213 is idle — see ``ServerRole._maybe_prune``.
214 """
215
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
216 + # Fetch rounds keyed by wire round_seq. A round is created by its
217 + # first supply or its fetch and removed at its terminal (finalize /
218 + # failure / abort).
219 + outbound: dict[int, _OutboundRequestState] = field(default_factory=dict)
220 # Raw inbound LookupMsgs not yet processed against the ParentManager.
221 pending_lookups: list[_PendingLookup] = field(default_factory=list)
222 # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by
223 # the (globally unique) lookup_id and re-polled each serve.
224 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
225 # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync
226 # via _inflight_add / _inflight_pop so a non-empty set is an exact
227 # "has any inflight transfer" predicate and the abort drain can
@@ -256,9 +269,9 @@ class ServerRole:
269 # ``parent.on_request_finished`` in ``serve_external_requests``.
270 self._finished_lookup_ctxs: list[ReqContext] = []
271 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()
272 + # Parked aborts awaiting drain, keyed by (kv_request_id, round)
273 + # with the abort start time.
274 + self._pending_aborts: dict[tuple[str, int], float] = {}
275
276 # ------------------------------------------------------------------
277 # State helpers
@@ -277,11 +290,11 @@ class ServerRole:
290 st = self._requests.get(kv_request_id)
291 if (
292 st is not None
280 - and st.outbound is None
293 + and not st.outbound
294 and not st.inflight_tids
295 and not st.lookups
296 and not st.pending_lookups
284 - and st.abort_started_at is None
297 + and not any(kv == kv_request_id for kv, _ in self._pending_aborts)
298 ):
299 del self._requests[kv_request_id]
300
@@ -295,87 +308,100 @@ class ServerRole:
308 keys: Sequence[OffloadKey],
309 block_ids: Sequence[int],
310 job_id: JobId,
311 + round_seq: int = 0,
312 + *,
313 + from_lookup: bool = False,
314 ) -> None:
299 - """New blocks stored locally — match against pending fetch demand."""
315 + """New blocks stored locally — match within their fetch round.
316 +
317 + Lookup pins carry the round they were probed under; PD
318 + submit_store batches share PD's single round 0.
319 + """
320 self._store_jobs[job_id] = time.monotonic()
321 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)
322 + rnd = st.outbound.get(round_seq)
323 + if rnd is None:
324 + rnd = st.outbound[round_seq] = _OutboundRequestState()
325 + if from_lookup:
326 + rnd.lookup_supplied = True
327 + result = rnd.add_stored_blocks(keys, block_ids, job_id)
328 + if result.local_idxs and rnd.demand_received:
329 + self._submit_transfer(kv_request_id, result, rnd, round_seq)
330
331 def on_fetch(
332 self,
333 kv_request_id: str,
334 keys: Sequence[OffloadKey],
335 block_indexes: Sequence[int],
336 + round_seq: int = 0,
337 ) -> None:
338 """Handle a FetchMsg from the peer.
339
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.
340 + A non-empty fetch binds and closes its round, leaving lookup
341 + state alone (the next round's LookupMsg may already be in
342 + flight). The terminal empty fetch closes the id: parked lookups
343 + are popped and every remaining round drained. A second fetch for
344 + a round already holding demand raises ValueError
345 + (protocol-error disconnect).
346 """
347 logger.debug(
341 - "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d",
348 + "P2PSession %s: fetch RECEIVED kv_request_id=%s round=%s blocks=%d",
349 self._peer_id,
350 kv_request_id,
351 + round_seq,
352 len(keys),
353 )
354 st = self._requests.get(kv_request_id)
347 - existing = st.outbound if st is not None else None
355 + existing = st.outbound.get(round_seq) if st is not None else None
356 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}")
357 + raise ValueError(
358 + f"duplicate fetch for kv_request_id={kv_request_id} round={round_seq}"
359 + )
360 st = self._get_or_create_request(kv_request_id)
354 - if st.outbound is None:
355 - st.outbound = _OutboundRequestState()
356 - req = st.outbound
361 + req = st.outbound.get(round_seq)
362 + if req is None:
363 + req = st.outbound[round_seq] = _OutboundRequestState()
364 result = req.add_fetch_demand(keys, block_indexes)
365 + if not keys:
366 + # Terminal empty fetch: close the lookup phase and drain
367 + # every round with no TransferDoneMsg (nothing waits on it).
368 + self._finish_inbound_lookups(kv_request_id)
369 + for key in list(st.outbound):
370 + self._finalize_outbound(kv_request_id, key, send_done=False)
371 + return
372 + if req.lookup_supplied and req.demanded:
373 + # A symmetric round's supply always precedes its fetch, so
374 + # unmatched demand is unservable — fail now, not at the load
375 + # timeout. PD rounds keep parking demand for stores that
376 + # arrive later.
377 + logger.warning(
378 + "P2PSession %s: fetch kv_request_id=%s round=%s demanded %d "
379 + "blocks but %d have no pinned supply; failing fetch "
380 + "immediately",
381 + self._peer_id,
382 + kv_request_id,
383 + round_seq,
384 + len(keys),
385 + len(req.demanded),
386 + )
387 + self._finalize_outbound(kv_request_id, round_seq, success=False)
388 + return
389 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)
390 + self._submit_transfer(kv_request_id, result, req, round_seq)
391 # Prefiller-first mode: finish_request may have run before
392 # fetch arrived. If so, finalize once we know what was
393 # demanded — fully satisfied → success, else early-fail.
369 - if req.finishing and not self._has_inflight_for(kv_request_id):
370 - self._finalize_outbound(kv_request_id)
394 + if req.finishing and req.inflight == 0:
395 + self._finalize_outbound(kv_request_id, round_seq)
396
372 - def on_abort_fetch(self, kv_request_id: str) -> None:
373 - """Handle an AbortFetchMsg from the peer."""
397 + def on_abort_fetch(self, kv_request_id: str, round_seq: int = 0) -> None:
398 + """Handle an AbortFetchMsg from the peer, cancelling one round."""
399 # Abort for an unknown id may be a benign race/duplicate or a
400 # real protocol violation; we don't track completed ids, so warn.
401 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):
402 + if (st is None or not st.outbound) and not self._has_inflight_for(
403 + kv_request_id
404 + ):
405 logger.warning(
406 "P2PSession %s: abort_fetch for unknown kv_request_id=%s "
407 "(no outbound or inflight state); benign race or stale",
@@ -385,16 +411,15 @@ class ServerRole:
411 # Idempotent: receiving AbortFetchMsg again before we've sent the
412 # ack just triggers another drain attempt without resetting the
413 # deadline.
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)
414 + self._get_or_create_request(kv_request_id)
415 + self._pending_aborts.setdefault((kv_request_id, round_seq), time.monotonic())
416 + self._drain_abort(kv_request_id, round_seq)
417
418 def on_lookup(
419 self,
420 kv_request_id: str,
421 keys: Sequence[OffloadKey],
422 + round_seq: int = 0,
423 ) -> None:
424 """Enqueue a LookupMsg from a symmetric-P2P consumer.
425
@@ -406,13 +431,18 @@ class ServerRole:
431 parent calls are valid.
432 """
433 logger.debug(
409 - "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d",
434 + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s round=%s keys=%d",
435 self._peer_id,
436 kv_request_id,
437 + round_seq,
438 len(keys),
439 )
440 self._get_or_create_request(kv_request_id).pending_lookups.append(
415 - _PendingLookup(keys=list(keys), enqueued_at=time.monotonic())
441 + _PendingLookup(
442 + keys=list(keys),
443 + enqueued_at=time.monotonic(),
444 + round_seq=round_seq,
445 + )
446 )
447 self._serve_pending.add(kv_request_id)
448
@@ -434,7 +464,7 @@ class ServerRole:
464 st.pending_lookups = []
465 for pl in pending:
466 self._process_inbound_lookup(
437 - kv_request_id, pl.keys, pl.enqueued_at, parent
467 + kv_request_id, pl.keys, pl.enqueued_at, pl.round_seq, parent
468 )
469 self._resolve_pending_lookups(kv_request_id, parent)
470 st = self._requests.get(kv_request_id)
@@ -481,9 +511,7 @@ class ServerRole:
511 else:
512 lookup.pending.add(h)
513 if new_hits:
484 - self._pin_and_register_hits(
485 - lookup.kv_request_id, new_hits, lookup.ctx, parent
486 - )
514 + self._pin_and_register_hits(lookup, new_hits, parent)
515 return new_hits
516
517 def _process_inbound_lookup(
@@ -491,6 +519,7 @@ class ServerRole:
519 kv_request_id: str,
520 keys: list[OffloadKey],
521 enqueued_at: float,
522 + round_seq: int,
523 parent: ParentManager,
524 ) -> None:
525 """Resolve one enqueued LookupMsg against ``parent``.
@@ -515,6 +544,7 @@ class ServerRole:
544 lookup_id=lookup_id,
545 kv_request_id=kv_request_id,
546 ctx=ctx,
547 + round_seq=round_seq,
548 keys=list(keys),
549 deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S,
550 )
@@ -547,25 +577,26 @@ class ServerRole:
577
578 def _pin_and_register_hits(
579 self,
550 - kv_request_id: str,
580 + lookup: _ActiveLookup,
581 keys: list[OffloadKey],
552 - ctx: ReqContext,
582 parent: ParentManager,
583 ) -> None:
555 - """Pin primary slots for HIT keys and feed them into the
556 - existing ``add_stored_blocks`` matching path.
584 + """Pin primary slots for HIT keys and park them as the lookup's
585 + round supply via ``add_stored_blocks``.
586
587 Caller has already confirmed every key is HIT (single-threaded
588 scheduler ⇒ no eviction race), so the JobMetadata returned by
589 ``parent.create_store_job`` carries parallel ``keys``/``block_ids``
590 of length ``len(keys)``.
591 """
563 - meta = parent.create_store_job(keys, ctx)
592 + meta = parent.create_store_job(keys, lookup.ctx)
593 self.add_stored_blocks(
565 - kv_request_id,
594 + lookup.kv_request_id,
595 list(meta.keys),
596 list(meta.block_ids),
597 meta.job_id,
598 + round_seq=lookup.round_seq,
599 + from_lookup=True,
600 )
601
602 def _resolve_pending_lookups(
@@ -650,9 +681,8 @@ class ServerRole:
681
682 Called on the two events that mean "no more lookup traffic for
683 ``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.
684 + empty FetchMsg from the peer and a local ``finish``. Whichever
685 + fires second is a no-op.
686 """
687 st = self._requests.get(kv_request_id)
688 if st is None:
@@ -688,20 +718,15 @@ class ServerRole:
718 self._finish_inbound_lookups(kv_request_id)
719
720 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
695 - if not req.demand_received:
696 - return
697 - if self._has_inflight_for(kv_request_id):
721 + if st is None:
722 return
699 - # Remaining > 0 here: if it had hit 0, the poll-done success
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)
723 + for key, req in list(st.outbound.items()):
724 + req.finishing = True
725 + if not req.demand_received or req.inflight:
726 + # No demand yet (prefiller-first): on_fetch finalizes via
727 + # `finishing`. Inflight: the last completion finalizes.
728 + continue
729 + self._finalize_outbound(kv_request_id, key)
730
731 def collect_results(self) -> list[StoreResult]:
732 """Drain timeouts, deferred results, and transport completions.
@@ -741,20 +766,24 @@ class ServerRole:
766 )
767 continue
768 results.extend(self._settle_xfer_jobs(xfer, success=True))
769 + rnd = xfer.round
770 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, (
771 + if st is not None and st.outbound.get(xfer.round_key) is rnd:
772 + rnd.remaining -= xfer.block_count
773 + assert rnd.remaining >= 0, (
774 f"remaining went negative for kv_request_id={xfer.kv_request_id}"
775 )
751 - if req.remaining == 0:
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)
776 + if rnd.remaining == 0:
777 + self._finalize_outbound(
778 + xfer.kv_request_id, xfer.round_key, success=True
779 + )
780 + elif rnd.finishing and rnd.inflight == 0:
781 + self._finalize_outbound(
782 + xfer.kv_request_id, xfer.round_key, success=False
783 + )
784 self._maybe_prune(xfer.kv_request_id)
785
757 - failed_kv_request_ids: set[str] | None = None
786 + failed_rounds: list[tuple[str, _OutboundRequestState]] | None = None
787 for tid in poll_result.failed:
788 xfer = self._inflight_pop(tid)
789 if xfer is None:
@@ -767,35 +796,36 @@ class ServerRole:
796 tid,
797 )
798 continue
770 - if failed_kv_request_ids is None:
771 - failed_kv_request_ids = set()
772 - failed_kv_request_ids.add(xfer.kv_request_id)
799 results.extend(self._settle_xfer_jobs(xfer, success=False))
800 + rnd = xfer.round
801 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:
802 + if st is not None and st.outbound.get(xfer.round_key) is rnd:
803 + del st.outbound[xfer.round_key]
804 + if failed_rounds is None:
805 + failed_rounds = []
806 + failed_rounds.append((xfer.kv_request_id, rnd))
807 self._send(
808 {
809 TYPE_KEY: TransferDoneMsg.TYPE,
810 TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id,
811 TransferDoneMsg.SUCCESS: False,
812 + TransferDoneMsg.ROUND_SEQ: xfer.round_key,
813 }
814 )
815 self._maybe_prune(xfer.kv_request_id)
816
788 - # Cancel other inflight for the same failed kv_request_ids
789 - if failed_kv_request_ids:
790 - ids_to_cancel = [
791 - tid
792 - for tid, xfer in self._inflight.items()
793 - if xfer.kv_request_id in failed_kv_request_ids
794 - ]
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:
817 + # Cancel each failed round's other inflight and fail its
818 + # remaining store jobs — nothing else will settle them.
819 + if failed_rounds:
820 + for kv_request_id, rnd in failed_rounds:
821 + ids_to_cancel = [
822 + tid for tid, x in self._inflight.items() if x.round is rnd
823 + ]
824 + for tid in ids_to_cancel:
825 + self._inflight_pop(tid)
826 + if ids_to_cancel:
827 + self._transport.cancel(ids_to_cancel)
828 + results.extend(self._fail_round_jobs(rnd))
829 self._maybe_prune(kv_request_id)
830
831 return results
@@ -811,8 +841,8 @@ class ServerRole:
841
842 def drain_pending_aborts(self) -> None:
843 """Re-attempt every parked abort once per poll tick."""
814 - for kv_request_id in list(self._parked_aborts):
815 - self._drain_abort(kv_request_id)
844 + for kv_request_id, round_seq in list(self._pending_aborts):
845 + self._drain_abort(kv_request_id, round_seq)
846
847 def close(self) -> tuple[list[int], list[ReqContext]]:
848 """Tear down. Cancels inflight.
@@ -839,7 +869,7 @@ class ServerRole:
869 failed_serves.extend(self._finished_lookup_ctxs)
870 self._requests.clear()
871 self._serve_pending.clear()
842 - self._parked_aborts.clear()
872 + self._pending_aborts.clear()
873 self._finished_lookup_ctxs.clear()
874 return failed_stores, failed_serves
875
@@ -870,6 +900,8 @@ class ServerRole:
900 xfer = self._inflight.pop(tid, None)
901 if xfer is None:
902 return None
903 + xfer.round.inflight -= 1
904 + assert xfer.round.inflight >= 0
905 st = self._requests.get(xfer.kv_request_id)
906 if st is not None:
907 st.inflight_tids.discard(tid)
@@ -880,80 +912,108 @@ class ServerRole:
912 ) -> list[StoreResult]:
913 """Emit StoreResults for a completed transfer's store jobs.
914
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.
915 + Pops each attached job from ``_store_jobs`` and clears it from
916 + its round's pending set. A job already popped (via timeout,
917 + cancel, etc.) is skipped so we never double-emit a contradictory
918 + result.
919 """
920 results: list[StoreResult] = []
888 - st = self._requests.get(xfer.kv_request_id)
889 - req = st.outbound if st is not None else None
921 for job_id in xfer.job_ids:
922 if self._store_jobs.pop(job_id, None) is None:
923 continue
924 results.append(StoreResult(job_id=job_id, success=success))
894 - if req is not None:
895 - req.pending_job_ids.discard(job_id)
925 + xfer.round.pending_job_ids.discard(job_id)
926 return results
927
928 # ------------------------------------------------------------------
929 # Internal — finalize / abort drain
930 # ------------------------------------------------------------------
931
932 + def _fail_round_jobs(self, rnd: _OutboundRequestState) -> list[StoreResult]:
933 + """Fail a terminated round's still-pending store jobs (idempotent)."""
934 + results: list[StoreResult] = []
935 + for job_id in rnd.pending_job_ids:
936 + if self._store_jobs.pop(job_id, None) is None:
937 + continue
938 + results.append(StoreResult(job_id=job_id, success=False))
939 + rnd.pending_job_ids.clear()
940 + return results
941 +
942 def _finalize_outbound(
943 self,
944 kv_request_id: str,
945 + round_key: int,
946 success: bool | None = None,
947 + send_done: bool = True,
948 ) -> None:
907 - """Pop the outbound state and emit terminal results.
908 -
909 - Called when no further work will happen for this kv_request_id
910 - on the server side: either request_finish has fired and there
911 - are no inflight transfers, or the last inflight just completed
912 - while finishing.
949 + """Pop one round and emit its terminal results.
950
951 If ``success`` is None, derive it from ``req.remaining == 0``.
915 - The same flag is used for both the peer's TransferDoneMsg and
916 - the StoreResult(s) emitted for any leftover pending job_ids.
952 + ``send_done=False`` skips the TransferDoneMsg (terminal empty
953 + fetch). Other rounds of the id are untouched.
954 """
955 st = self._requests[kv_request_id]
919 - assert st.outbound is not None
920 - req = st.outbound
921 - st.outbound = None
956 + req = st.outbound.pop(round_key)
957 if success is None:
923 - success = req.remaining == 0
924 - for job_id in req.pending_job_ids:
925 - self._store_jobs.pop(job_id, None)
926 - self._pending_store_results.append(
927 - StoreResult(job_id=job_id, success=success)
928 - )
929 - self._send(
930 - {
931 - TYPE_KEY: TransferDoneMsg.TYPE,
932 - TransferDoneMsg.KV_REQUEST_ID: kv_request_id,
933 - TransferDoneMsg.SUCCESS: success,
934 - }
958 + success = req.demand_received and req.remaining == 0
959 + settled = self._fail_round_jobs(req) if not success else None
960 + if settled is not None:
961 + self._pending_store_results.extend(settled)
962 + else:
963 + for job_id in req.pending_job_ids:
964 + if self._store_jobs.pop(job_id, None) is None:
965 + continue
966 + self._pending_store_results.append(
967 + StoreResult(job_id=job_id, success=True)
968 + )
969 + req.pending_job_ids.clear()
970 + logger.debug(
971 + "P2PSession %s: finalize kv_request_id=%s round=%s success=%s "
972 + "remaining=%d leftover_available=%d send_done=%s",
973 + self._peer_id,
974 + kv_request_id,
975 + round_key,
976 + success,
977 + req.remaining,
978 + len(req.available),
979 + send_done,
980 )
981 + if send_done and req.demand_received:
982 + self._send(
983 + {
984 + TYPE_KEY: TransferDoneMsg.TYPE,
985 + TransferDoneMsg.KV_REQUEST_ID: kv_request_id,
986 + TransferDoneMsg.SUCCESS: success,
987 + TransferDoneMsg.ROUND_SEQ: round_key,
988 + }
989 + )
990 self._maybe_prune(kv_request_id)
991
938 - def _drain_abort(self, kv_request_id: str) -> None:
992 + def _drain_abort(self, kv_request_id: str, round_seq: int) -> None:
993 """One drain attempt for a pending abort.
994
941 - Stops accepting more blocks for ``kv_request_id``, then asks the
942 - transport to cancel any matching inflight transfers in
943 - ``mode="wait"``. Sends ``AbortAckMsg`` once nothing remains
944 - inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to
945 - ``mode="immediate"`` and acks anyway.
995 + Detaches the aborted round, then asks the transport to cancel its
996 + inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` once
997 + nothing remains inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S``
998 + falls back to ``mode="immediate"`` and acks anyway.
999 """
1000 st = self._requests[kv_request_id]
948 - st.outbound = None
949 - ids = list(st.inflight_tids)
1001 + rnd = st.outbound.pop(round_seq, None)
1002 + if rnd is not None:
1003 + # Its transfers are being cancelled; fail its jobs now
1004 + # instead of leaking them to the store timeout.
1005 + self._pending_store_results.extend(self._fail_round_jobs(rnd))
1006 + ids = [
1007 + tid
1008 + for tid, x in self._inflight.items()
1009 + if x.kv_request_id == kv_request_id and x.round_key == round_seq
1010 + ]
1011 if not ids:
951 - self._finalize_abort(kv_request_id)
1012 + self._finalize_abort(kv_request_id, round_seq)
1013 return
1014
954 - assert st.abort_started_at is not None
955 - expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S
956 - if expired:
1015 + started_at = self._pending_aborts[(kv_request_id, round_seq)]
1016 + if time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S:
1017 for tid in ids:
1018 self._inflight_pop(tid)
1019 self._transport.cancel(ids, mode="immediate")
@@ -964,7 +1024,7 @@ class ServerRole:
1024 kv_request_id,
1025 len(ids),
1026 )
967 - self._finalize_abort(kv_request_id)
1027 + self._finalize_abort(kv_request_id, round_seq)
1028 return
1029
1030 still = self._transport.cancel(ids, mode="wait")
@@ -977,16 +1037,15 @@ class ServerRole:
1037 if tid not in still_set:
1038 self._inflight_pop(tid)
1039 if not still:
980 - self._finalize_abort(kv_request_id)
1040 + self._finalize_abort(kv_request_id, round_seq)
1041
982 - def _finalize_abort(self, kv_request_id: str) -> None:
983 - st = self._requests[kv_request_id]
984 - st.abort_started_at = None
985 - self._parked_aborts.discard(kv_request_id)
1042 + def _finalize_abort(self, kv_request_id: str, round_seq: int) -> None:
1043 + self._pending_aborts.pop((kv_request_id, round_seq), None)
1044 self._send(
1045 {
1046 TYPE_KEY: AbortAckMsg.TYPE,
1047 AbortAckMsg.KV_REQUEST_ID: kv_request_id,
1048 + AbortAckMsg.ROUND_SEQ: round_seq,
1049 }
1050 )
1051 self._maybe_prune(kv_request_id)
@@ -995,7 +1054,13 @@ class ServerRole:
1054 # Internal — transfers and store-job timeouts
1055 # ------------------------------------------------------------------
1056
998 - def _submit_transfer(self, kv_request_id: str, result: _MatchResult) -> None:
1057 + def _submit_transfer(
1058 + self,
1059 + kv_request_id: str,
1060 + result: _MatchResult,
1061 + rnd: _OutboundRequestState,
1062 + round_key: int,
1063 + ) -> None:
1064 logger.debug(
1065 "P2PSession %s: NIXL write_blocks CALL kv_request_id=%s "
1066 "local_idxs=%d remote_idxs=%d",
@@ -1016,12 +1081,15 @@ class ServerRole:
1081 transfer_id,
1082 len(result.local_idxs),
1083 )
1084 + rnd.inflight += 1
1085 self._inflight_add(
1086 transfer_id,
1087 _InflightXfer(
1088 kv_request_id=kv_request_id,
1089 block_count=len(result.local_idxs),
1090 job_ids=result.job_ids,
1091 + round=rnd,
1092 + round_key=round_key,
1093 ),
1094 )
1095 else:
@@ -1031,22 +1099,24 @@ class ServerRole:
1099 kv_request_id,
1100 len(result.local_idxs),
1101 )
1034 - # The matched blocks were popped from req.demanded /
1035 - # req.available, but no inflight will satisfy them, so
1036 - # remaining will never reach 0 on its own. Mark the
1037 - # request as finishing so the existing terminal paths
1038 - # clean up: if other inflight is in flight, the last one
1039 - # to drain will fire _finalize_outbound(success=False)
1040 - # via the elif branch in collect_results. If
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.
1102 + # The matched blocks were popped from rnd.demanded /
1103 + # rnd.available, but no inflight will satisfy them, so
1104 + # remaining will never reach 0 on its own. Mark the round
1105 + # as finishing so the existing terminal paths clean up: if
1106 + # other transfers of this round are in flight, the last one
1107 + # to drain will fire _finalize_outbound(success=False) via
1108 + # the elif branch in collect_results. If nothing else is in
1109 + # flight, finalize now so the peer and the local store jobs
1110 + # don't wait for finish_request or for _STORE_TIMEOUT_S /
1111 + # _LOAD_TIMEOUT_S.
1112 + rnd.finishing = True
1113 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):
1049 - self._finalize_outbound(kv_request_id, success=False)
1114 + if (
1115 + st is not None
1116 + and st.outbound.get(round_key) is rnd
1117 + and rnd.inflight == 0
1118 + ):
1119 + self._finalize_outbound(kv_request_id, round_key, success=False)
1120
1121 def _timeout_pending_store_jobs(self) -> list[StoreResult]:
1122 if not self._store_jobs:
vllm/v1/kv_offload/tiering/p2p/session/session.py
+12 -4
@@ -374,26 +374,34 @@ class P2PSession:
374 for bh in msg[FetchMsg.KEYS]
375 ]
376 block_indexes = msg[FetchMsg.BLOCK_INDEXES]
377 + round_seq = msg[FetchMsg.ROUND_SEQ]
378 # Run the server-role state machine inline as today —
379 # add_fetch_demand records demand against any blocks we've
380 # already seen in `available`. Report the kv_request_id so
381 # the manager (after poll() returns) can replay any parked
382 # submit_store batches; their add_stored_blocks calls hit
383 # the demand recorded here and submit transfers immediately.
383 - self._server.on_fetch(kv_request_id, keys, block_indexes)
384 + self._server.on_fetch(kv_request_id, keys, block_indexes, round_seq)
385 self._new_fetch_ids.append(kv_request_id)
386 elif msg_type == AbortFetchMsg.TYPE:
387 AbortFetchMsg.validate(msg)
387 - self._server.on_abort_fetch(msg[AbortFetchMsg.KV_REQUEST_ID])
388 + self._server.on_abort_fetch(
389 + msg[AbortFetchMsg.KV_REQUEST_ID],
390 + msg[AbortFetchMsg.ROUND_SEQ],
391 + )
392 elif msg_type == TransferDoneMsg.TYPE:
393 TransferDoneMsg.validate(msg)
394 self._client.on_transfer_done(
395 msg[TransferDoneMsg.KV_REQUEST_ID],
396 msg[TransferDoneMsg.SUCCESS],
397 + msg[TransferDoneMsg.ROUND_SEQ],
398 )
399 elif msg_type == AbortAckMsg.TYPE:
400 AbortAckMsg.validate(msg)
396 - self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID])
401 + self._client.on_abort_ack(
402 + msg[AbortAckMsg.KV_REQUEST_ID],
403 + msg[AbortAckMsg.ROUND_SEQ],
404 + )
405 elif msg_type == LookupMsg.TYPE:
406 LookupMsg.validate(msg)
407 kv_request_id = msg[LookupMsg.KV_REQUEST_ID]
@@ -401,7 +409,7 @@ class P2PSession:
409 OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))
410 for bh in msg[LookupMsg.KEYS]
411 ]
404 - self._server.on_lookup(kv_request_id, keys)
412 + self._server.on_lookup(kv_request_id, keys, msg[LookupMsg.ROUND_SEQ])
413 elif msg_type == LookupRespMsg.TYPE:
414 LookupRespMsg.validate(msg)
415 kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID]