[Bugfix][KV Offload][P2P] Preserve aborted loads until abort completion (#52571)
Signed-off-by: Uttam Kumar <uttam.kumar912@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Uttam Kumar committed
Aug 31, 2026 at 04:34 UTC
e0d27040ddcc5ac31cf01c5b04a7d764ccba656d
2 files changed
+97
-11
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+88
@@ -534,6 +534,94 @@ class TestClientFlows:
534
assert abort[TYPE_KEY] == AbortFetchMsg.TYPE
535
assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1"
536
537
+ def test_finish_preserves_active_load_until_terminal_result(self):
538
+ session, conn, _ = _make_session()
539
+ _activate(session, conn)
540
+ session.request_blocks(
541
+ job_id=7, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
542
+ )
543
+
544
+ session.finish_request("req-1")
545
+
546
+ load = _client_load(session, "req-1")
547
+ assert load.job_id == 7
548
+ assert load.aborted_at is not None
549
+ aborted_at = load.aborted_at
550
+ assert session._client.has_active_loads is True
551
+ session.finish_request("req-1")
552
+ assert _client_load(session, "req-1").aborted_at == aborted_at
553
+ aborts = [m for m in conn._sent if m[TYPE_KEY] == AbortFetchMsg.TYPE]
554
+ assert len(aborts) == 1
555
+ assert aborts[0][AbortFetchMsg.ROUND_SEQ] == 0
556
+
557
+ def test_finish_abort_ack_emits_one_failure(self):
558
+ session, conn, _ = _make_session()
559
+ _activate(session, conn)
560
+ session.request_blocks(
561
+ job_id=8, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
562
+ )
563
+ session.finish_request("req-1")
564
+
565
+ conn.enqueue(
566
+ {
567
+ TYPE_KEY: AbortAckMsg.TYPE,
568
+ AbortAckMsg.ROUND_SEQ: 0,
569
+ AbortAckMsg.KV_REQUEST_ID: "req-1",
570
+ }
571
+ )
572
+ loads = session.poll().loads
573
+ assert loads == [LoadResult(job_id=8, kv_request_id="req-1", success=False)]
574
+ assert session._client.has_active_loads is False
575
+ assert "req-1" not in session._client._requests
576
+ assert session.poll().loads == []
577
+
578
+ def test_finish_abort_timeout_emits_one_failure(self):
579
+ session, conn, _ = _make_session()
580
+ _activate(session, conn)
581
+ session.request_blocks(
582
+ job_id=9, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
583
+ )
584
+ session.finish_request("req-1")
585
+ _client_load(session, "req-1").aborted_at = (
586
+ time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
587
+ )
588
+
589
+ loads = session.poll().loads
590
+ assert loads == [LoadResult(job_id=9, kv_request_id="req-1", success=False)]
591
+ assert session._client.has_active_loads is False
592
+ assert "req-1" not in session._client._requests
593
+ assert session.poll().loads == []
594
+
595
+ def test_late_transfer_done_after_abort_ack_is_ignored(self):
596
+ session, conn, _ = _make_session()
597
+ _activate(session, conn)
598
+ session.request_blocks(
599
+ job_id=10, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
600
+ )
601
+ session.finish_request("req-1")
602
+
603
+ conn.enqueue(
604
+ {
605
+ TYPE_KEY: AbortAckMsg.TYPE,
606
+ AbortAckMsg.ROUND_SEQ: 0,
607
+ AbortAckMsg.KV_REQUEST_ID: "req-1",
608
+ }
609
+ )
610
+ assert session.poll().loads == [
611
+ LoadResult(job_id=10, kv_request_id="req-1", success=False)
612
+ ]
613
+
614
+ conn.enqueue(
615
+ {
616
+ TYPE_KEY: TransferDoneMsg.TYPE,
617
+ TransferDoneMsg.ROUND_SEQ: 0,
618
+ TransferDoneMsg.KV_REQUEST_ID: "req-1",
619
+ TransferDoneMsg.SUCCESS: True,
620
+ }
621
+ )
622
+ assert session.poll().loads == []
623
+ assert session._client.has_active_loads is False
624
+
625
def test_active_loads_work_list_tracks_in_flight(self):
626
"""collect_results / has_active_loads use the _active_loads work-list,
627
armed when a fetch is issued and discarded exactly when its load
vllm/v1/kv_offload/tiering/p2p/session/client.py
+9
-11
@@ -130,10 +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()
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``.
133
+ # kv_request_ids with at least one fetch in flight or aborting — the
134
+ # work-list collect_results walks for timeouts, and the
135
+ # has_active_loads predicate, instead of scanning every request.
136
+ # Kept in exact sync with ``st.loads``.
137
self._active_loads: set[str] = set()
138
self._completed_loads: list[LoadResult] = []
139
@@ -241,7 +241,8 @@ class ClientRole:
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.
244
+ Then drop all probe/lookup state. Keep aborted loads until their abort
245
+ acknowledgement or timeout produces a terminal result.
246
"""
247
st = self._requests.get(kv_request_id)
248
if st is None:
@@ -257,8 +258,7 @@ class ClientRole:
258
AbortFetchMsg.ROUND_SEQ: round_seq,
259
}
260
)
260
- st.loads.clear()
261
- self._active_loads.discard(kv_request_id)
261
+ load.aborted_at = time.monotonic()
262
if st.peer_lookup_open:
263
st.peer_lookup_open = False
264
self._send(
@@ -311,10 +311,8 @@ class ClientRole:
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 all "
316
- "nodes use the same prefix-cache hash seed (matching "
317
- "PYTHONHASHSEED, if set) and hash algorithm.",
314
+ "P2PSession %s: load request %s (job_id=%d) "
315
+ "completed with failure after abort",
316
self._peer_id,
317
kv_request_id,
318
load.job_id,