From e2aabc8b1521aad382b439e6be4e740d8e8211eb Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Wed, 14 Oct 2026 23:31:31 +0800 Subject: [PATCH] [Bugfix][KV Offload][P2P] Keep store pins and AbortAck until NIXL transfers stop (#59790) Signed-off-by: apex-mochen <2756823972@qq.com> Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .../tiering/p2p/test_data_transport.py | 89 +++- .../kv_offload/tiering/p2p/test_sessions.py | 408 ++++++++++++++++-- vllm/v1/kv_offload/tiering/p2p/data/base.py | 12 +- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 53 ++- vllm/v1/kv_offload/tiering/p2p/manager.py | 9 +- .../kv_offload/tiering/p2p/session/server.py | 232 +++++++--- 6 files changed, 695 insertions(+), 108 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index 60830a34bb..7219d691fd 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -14,6 +14,15 @@ import pytest from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport + +class nixlRemoteDisconnectError(Exception): # noqa: N801 + """Stand-in for NIXL's binding exception of the same name.""" + + +def _remote_disconnect() -> Exception: + return nixlRemoteDisconnectError("NIXL_ERR_REMOTE_DISCONNECT") + + # --------------------------------------------------------------------------- # DataTransport base class tests # --------------------------------------------------------------------------- @@ -286,25 +295,101 @@ class TestNixlTransportWithMockedAgent: assert transport.cancel([999, 1000]) == [] assert transport.cancel([999, 1000], mode="wait") == [] - def test_cancel_wait_release_succeeds(self): - """wait-mode cancel that succeeds pops the entry and returns [].""" + @pytest.mark.parametrize("state", ["DONE", "ERR"]) + def test_cancel_wait_release_succeeds(self, state): + """wait-mode cancel of a stopped transfer pops it and returns [].""" transport = self._make_transport() transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) tid = transport.write_blocks("peer:1", [0], [1]) assert tid in transport._inflight + transport._agent.check_xfer_state.return_value = state result = transport.cancel([tid], mode="wait") assert result == [] assert tid not in transport._inflight transport._agent.release_xfer_handle.assert_called_once() + @pytest.mark.parametrize("state", ["PROC", "PEND"]) + def test_cancel_wait_keeps_running_transfer(self, state): + """Releasing a running NIXL WRITE does not stop it, so wait-mode + cancel must not release it or report it as cancelled.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + + tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.return_value = state + + assert transport.cancel([tid], mode="wait") == [tid] + assert tid in transport._inflight + transport._agent.release_xfer_handle.assert_not_called() + + def test_cancel_wait_state_check_raises(self): + """A failed state check is not proof the transfer stopped.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + + tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.side_effect = RuntimeError("NOT_FOUND") + + assert transport.cancel([tid], mode="wait") == [tid] + assert tid in transport._inflight + transport._agent.release_xfer_handle.assert_not_called() + + def test_cancel_wait_remote_disconnect_is_terminal(self): + """A dead peer can no longer be written to: wait-mode cancel drops + the transfer even if its handle can no longer be released.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + + tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.side_effect = _remote_disconnect() + transport._agent.release_xfer_handle.side_effect = RuntimeError("gone") + + assert transport.cancel([tid], mode="wait") == [] + assert tid not in transport._inflight + assert transport.poll() == PollResult(done=(), failed=()) + + @pytest.mark.parametrize( + "exc", + [ + _remote_disconnect(), + # Recognized by status code even if the binding class is renamed. + RuntimeError("getXferStatus: NIXL_ERR_REMOTE_DISCONNECT"), + ], + ) + def test_poll_reports_remote_disconnect_as_failed(self, exc): + """check_xfer_state raises instead of returning ERR once the peer + disconnected; poll must surface that as a failure, not keep the + transfer inflight forever.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + + tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.side_effect = exc + + result = transport.poll() + assert result.done == () + assert tid in result.failed + assert tid not in transport._inflight + + def test_poll_keeps_transfer_on_other_state_check_errors(self): + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + + tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.side_effect = RuntimeError("NOT_FOUND") + + assert transport.poll() == PollResult(done=(), failed=()) + assert tid in transport._inflight + def test_cancel_wait_release_raises(self): """wait-mode cancel keeps the entry and returns the tid on raise.""" transport = self._make_transport() transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) tid = transport.write_blocks("peer:1", [0], [1]) + transport._agent.check_xfer_state.return_value = "DONE" transport._agent.release_xfer_handle.side_effect = RuntimeError( "NIXL_ERR_REPOST_ACTIVE" ) diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index adb3eada22..63abc642b0 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -15,10 +15,12 @@ from __future__ import annotations import time from collections.abc import Callable, Sequence from concurrent.futures import Future +from types import SimpleNamespace import numpy as np import pytest +import vllm.v1.kv_offload.tiering.p2p.session.server as server_mod from vllm.v1.kv_offload.base import ( LookupResult, OffloadKey, @@ -1632,9 +1634,11 @@ class TestServerFlows: assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight - def test_abort_fetch_force_cancels_after_timeout(self): - """If wait-mode never drains, the deadline forces immediate - cancel and an ack is still sent.""" + def test_abort_fetch_never_force_cancels_after_timeout(self): + """Past the drain deadline a transfer that is still running is not + force-released and no ack is sent: releasing the handle or removing + the peer does not stop a NIXL WRITE, so the ack would let the peer + reuse blocks that are still being written.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 @@ -1658,21 +1662,28 @@ class TestServerFlows: ) session.poll() assert _srv_abort_started(session, "req-1") is not None + peers = dict(transport._remote_peers) # Backdate past the drain deadline. session._server._pending_aborts[("req-1", 0)] = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) - # Even if the transport still claims it can't cancel, the - # session must force-pop and ack. transport._cancel_calls.clear() - session.poll() + for _ in range(3): + session.poll() + assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) + assert tid in session._server._inflight + assert all(mode == "wait" for _, mode in transport._cancel_calls) + assert transport._remote_peers == peers - ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) - assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" + # Once NIXL reports the transfer terminal, the ack follows. + transport._cancel_still_inflight.discard(tid) + transport._poll_done.append(tid) + session.poll() + acks = [m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE] + assert len(acks) == 1 assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight - assert ([tid], "immediate") in transport._cancel_calls def test_abort_fetch_idempotent_while_draining(self): """Receiving AbortFetchMsg twice for the same kv_request_id @@ -1731,10 +1742,12 @@ class TestServerFlows: stores = session.poll().stores assert StoreResult(job_id=1, success=False) in stores - def test_store_timeout_then_late_completion_no_duplicate(self): - """A job timed out by _timeout_pending_store_jobs must not also - emit a contradictory StoreResult(success=True) when the transport - later reports the same transfer as done.""" + @pytest.mark.parametrize("cancel_pending", [False, True]) + def test_store_timeout_cancels_stuck_transfer_after_confirmation( + self, cancel_pending, monkeypatch + ): + # Retry on every poll; pacing is covered by the *_paced tests. + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1749,25 +1762,143 @@ class TestServerFlows: ) session.poll() tid = next(iter(transport._transfers)) - - # Backdate the store job so the next poll times it out. session._server._store_jobs[1] = time.monotonic() - 60.0 - stores = session.poll().stores - assert StoreResult(job_id=1, success=False) in stores - assert StoreResult(job_id=1, success=True) not in stores + if cancel_pending: + transport._cancel_still_inflight.add(tid) + for _ in range(3): + assert session.poll().stores == [] + assert 1 in session._server._store_jobs + assert tid in session._server._inflight + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + transport._cancel_still_inflight.clear() + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert transport._cancel_calls + assert all(mode == "wait" for _, mode in transport._cancel_calls) + assert tid not in session._server._inflight + assert tid not in transport._transfers + terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE] + assert len(terminal) == 1 + assert terminal[0][TransferDoneMsg.SUCCESS] is False + assert session.poll().stores == [] - # Transport later reports the same transfer as done — must not - # emit a second (contradictory) StoreResult for job_id=1. - transport._poll_done.append(tid) - stores = session.poll().stores - assert all(s.job_id != 1 for s in stores), ( - f"unexpected duplicate StoreResult after timeout: {stores}" + @pytest.mark.parametrize("success", [True, False]) + def test_store_timeout_keeps_pins_until_cancel_or_late_terminal(self, success): + session, conn, transport = _make_session() + _activate(session, conn) + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) + session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2) + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"k1"], + FetchMsg.BLOCK_INDEXES: [5], + } ) + session.poll() + tid = next(iter(transport._transfers)) + transport._cancel_still_inflight.add(tid) - def test_store_timeout_then_late_failure_no_duplicate(self): - """Symmetric guard: a timed-out job must not also emit a second - StoreResult(success=False) when the transport later reports the - same transfer as failed.""" + for jid in (1, 2): + session._server._store_jobs[jid] = time.monotonic() - 60.0 + # A StoreResult releases the primary slot. Only the job still + # waiting for a fetch may time out while the transfer is active. + assert session.poll().stores == [StoreResult(job_id=2, success=False)] + assert session.poll().stores == [] + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + + completed = transport._poll_done if success else transport._poll_failed + completed.append(tid) + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE] + assert len(terminal) == 1 + assert terminal[0][TransferDoneMsg.SUCCESS] is False + assert session.poll().stores == [] + + @pytest.mark.parametrize("success", [True, False, None]) + def test_store_timeout_waits_for_other_transfers_in_round( + self, success, monkeypatch + ): + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) + session, conn, transport = _make_session() + _activate(session, conn) + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"k1", b"k2"], + FetchMsg.BLOCK_INDEXES: [5, 6], + } + ) + session.poll() + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) + session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2) + first, second = list(transport._transfers) + session._server._store_jobs[1] = time.monotonic() - 60.0 + transport._cancel_still_inflight.add(second) + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert first not in session._server._inflight + assert second in session._server._inflight + assert 2 in session._server._store_jobs + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + # Additional supply must not restart the timed-out round. + session.add_stored_blocks("req-1", [b"k3"], [2], job_id=3) + assert session.poll().stores == [StoreResult(job_id=3, success=False)] + assert set(session._server._inflight) == {second} + if success is None: + transport._cancel_still_inflight.clear() + else: + completed = transport._poll_done if success else transport._poll_failed + completed.append(second) + assert session.poll().stores == [StoreResult(job_id=2, success=False)] + terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE] + assert len(terminal) == 1 + assert terminal[0][TransferDoneMsg.SUCCESS] is False + assert session.poll().stores == [] + + def test_failed_transfer_keeps_sibling_pins_until_terminal(self, monkeypatch): + """A failed transfer fails its round, but a sibling transfer that may + still be writing is not force-cancelled: its pin and the round's + TransferDoneMsg wait until it is terminal.""" + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) + session, conn, transport = _make_session() + _activate(session, conn) + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"k1", b"k2"], + FetchMsg.BLOCK_INDEXES: [5, 6], + } + ) + session.poll() + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) + session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2) + first, second = list(transport._transfers) + transport._cancel_still_inflight.add(second) + + transport._poll_failed.append(first) + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + for _ in range(3): + assert session.poll().stores == [] + assert second in session._server._inflight + assert 2 in session._server._store_jobs + assert transport._cancel_calls + assert all(mode == "wait" for _, mode in transport._cancel_calls) + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + + transport._cancel_still_inflight.clear() + assert session.poll().stores == [StoreResult(job_id=2, success=False)] + terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE] + assert len(terminal) == 1 + assert terminal[0][TransferDoneMsg.SUCCESS] is False + assert session.poll().stores == [] + + def test_abort_during_store_timeout_keeps_unconfirmed_pins(self, monkeypatch): + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1782,18 +1913,219 @@ class TestServerFlows: ) session.poll() tid = next(iter(transport._transfers)) - + transport._cancel_still_inflight.add(tid) session._server._store_jobs[1] = time.monotonic() - 60.0 - stores = session.poll().stores - assert [s for s in stores if s.job_id == 1] == [ - StoreResult(job_id=1, success=False) - ] - - transport._poll_failed.append(tid) - stores = session.poll().stores - assert all(s.job_id != 1 for s in stores), ( - f"unexpected duplicate StoreResult after timeout: {stores}" + assert session.poll().stores == [] + conn.enqueue( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: "req-1", + AbortFetchMsg.ROUND_SEQ: 0, + } ) + assert session.poll().stores == [] + session._server._pending_aborts[("req-1", 0)] = time.monotonic() - 60.0 + assert session.poll().stores == [] + assert 1 in session._server._store_jobs + assert tid in session._server._inflight + assert all(mode == "wait" for _, mode in transport._cancel_calls) + assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) + transport._cancel_still_inflight.clear() + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + assert session.poll().stores == [] + + @staticmethod + def _fake_clock(monkeypatch) -> list[float]: + clock = [1000.0] + monkeypatch.setattr( + server_mod, "time", SimpleNamespace(monotonic=lambda: clock[0]) + ) + return clock + + @staticmethod + def _submit_unconfirmed_transfer(session, conn, transport) -> int: + """Start one transfer for job 1 whose wait-cancel never confirms.""" + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"k1"], + FetchMsg.BLOCK_INDEXES: [5], + } + ) + session.poll() + tid = next(iter(transport._transfers)) + transport._cancel_still_inflight.add(tid) + return tid + + @staticmethod + def _wait_cancels(transport) -> int: + return sum(1 for _, mode in transport._cancel_calls if mode == "wait") + + def test_store_timeout_cancel_retry_is_paced(self, monkeypatch): + clock = self._fake_clock(monkeypatch) + session, conn, transport = _make_session() + _activate(session, conn) + tid = self._submit_unconfirmed_transfer(session, conn, transport) + session._server._store_jobs[1] = clock[0] - 60.0 + + assert session.poll().stores == [] + assert self._wait_cancels(transport) == 1 + start = clock[0] + for dt in (0.0, 0.05, 0.099): + clock[0] = start + dt + for _ in range(5): + assert session.poll().stores == [] + assert self._wait_cancels(transport) == 1 + + clock[0] = start + server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3 + assert session.poll().stores == [] + assert self._wait_cancels(transport) == 2 + + transport._cancel_still_inflight.clear() + clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3 + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert self._wait_cancels(transport) == 3 + assert tid not in session._server._inflight + + def test_abort_drain_cancel_retry_is_paced(self, monkeypatch): + clock = self._fake_clock(monkeypatch) + session, conn, transport = _make_session() + _activate(session, conn) + self._submit_unconfirmed_transfer(session, conn, transport) + conn.enqueue( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: "req-1", + AbortFetchMsg.ROUND_SEQ: 0, + } + ) + # The first drain attempt is immediate. + session.poll() + assert self._wait_cancels(transport) == 1 + for _ in range(5): + session.poll() + assert self._wait_cancels(transport) == 1 + + clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3 + session.poll() + assert self._wait_cancels(transport) == 2 + assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) + + transport._cancel_still_inflight.clear() + clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3 + session.poll() + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + + def test_unconfirmed_store_cancel_never_acks_or_tears_down(self, monkeypatch): + """A cancellation that is never confirmed must not release the pin, + send AbortAck, force-cancel, or remove the peer, however long it + stays pending. AbortAck follows the release once it is confirmed.""" + clock = self._fake_clock(monkeypatch) + session, conn, transport = _make_session() + _activate(session, conn) + peers = dict(transport._remote_peers) + tid = self._submit_unconfirmed_transfer(session, conn, transport) + session._server._store_jobs[1] = clock[0] - 60.0 + assert session.poll().stores == [] + + abort = { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: "req-1", + AbortFetchMsg.ROUND_SEQ: 0, + } + started_at = clock[0] + conn.enqueue(dict(abort)) + assert session.poll().stores == [] + # A duplicate abort neither resets the drain start nor adds a cancel. + conn.enqueue(dict(abort)) + assert session.poll().stores == [] + assert session._server._pending_aborts[("req-1", 0)] == started_at + # The store sweep and the abort drain share the round's pacing. + assert self._wait_cancels(transport) == 1 + + ticks = 50 + for step in range(1, ticks + 1): + clock[0] = started_at + step * 10.0 + assert session.poll().stores == [] + assert clock[0] - started_at > _CANCEL_DRAIN_TIMEOUT_S + assert self._wait_cancels(transport) == 1 + ticks + assert all(mode == "wait" for _, mode in transport._cancel_calls) + assert 1 in session._server._store_jobs + assert tid in session._server._inflight + assert transport._remote_peers == peers + assert not any( + m[TYPE_KEY] in (AbortAckMsg.TYPE, TransferDoneMsg.TYPE) for m in conn._sent + ) + + transport._cancel_still_inflight.clear() + clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3 + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert 1 not in session._server._store_jobs + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + assert session.poll().stores == [] + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + + def test_abort_keeps_pins_only_of_jobs_being_written(self, monkeypatch): + """An abort fails the round's jobs at once, except those an inflight + transfer still reads; their pins and the ack wait for the transfer.""" + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) + session, conn, transport = _make_session() + _activate(session, conn) + tid = self._submit_unconfirmed_transfer(session, conn, transport) + # Job 2 supplies a block nobody demanded: no transfer reads it. + session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2) + conn.enqueue( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: "req-1", + AbortFetchMsg.ROUND_SEQ: 0, + } + ) + assert session.poll().stores == [StoreResult(job_id=2, success=False)] + for _ in range(3): + assert session.poll().stores == [] + assert 1 in session._server._store_jobs + assert tid in session._server._inflight + assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) + + transport._cancel_still_inflight.clear() + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + assert session.poll().stores == [] + + @pytest.mark.parametrize("success", [True, False]) + def test_abort_acks_after_late_terminal_without_transfer_done( + self, success, monkeypatch + ): + """A transfer that finishes on its own after the abort releases its + pin and lets the parked abort ack; the aborted round sends no + TransferDoneMsg.""" + monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0) + session, conn, transport = _make_session() + _activate(session, conn) + tid = self._submit_unconfirmed_transfer(session, conn, transport) + conn.enqueue( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: "req-1", + AbortFetchMsg.ROUND_SEQ: 0, + } + ) + assert session.poll().stores == [] + assert 1 in session._server._store_jobs + + completed = transport._poll_done if success else transport._poll_failed + completed.append(tid) + assert session.poll().stores == [StoreResult(job_id=1, success=False)] + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 + assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) + assert session.poll().stores == [] + assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1 def test_fetch_after_symmetric_supply_timed_out_reads_nothing(self): """A lookup pin whose store job timed out has been reported failed, diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index 574ba8955d..544ca8ebc1 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -296,12 +296,12 @@ class DataTransport(ABC): return []. Matches the legacy fire-and-forget behavior — the caller does not wait for the underlying transfer to drain. - "wait": attempt to release each handle. If the release - cannot complete because the transfer is still - PROC/PEND, the entry stays in the inflight set and - its id is included in the returned list. The - caller is expected to keep calling poll() until - every returned id surfaces in done/failed. + "wait": release only transfers that have stopped. A + transfer still PROC/PEND (or whose release fails) + stays in the inflight set and its id is included in + the returned list. The caller is expected to keep + calling poll() until every returned id surfaces in + done/failed. Returns: For mode="wait", the subset of *transfer_ids* still diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index c44069ceec..4e7951f320 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -29,6 +29,22 @@ logger = init_logger(__name__) _EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=()) +def _is_remote_disconnect(exc: Exception) -> bool: + """True if a NIXL state check failed because the peer disconnected. + + NIXL raises ``nixlRemoteDisconnectError`` (NIXL_ERR_REMOTE_DISCONNECT) + from ``check_xfer_state`` once the target endpoint is gone, e.g. after + the peer process died, instead of reporting an ``"ERR"`` state. Matched + by name and status code so the bindings need not be imported here; an + unrecognized error keeps the transfer inflight, the conservative side. + """ + return type( + exc + ).__name__ == "nixlRemoteDisconnectError" or "NIXL_ERR_REMOTE_DISCONNECT" in str( + exc + ) + + class _Inflight(NamedTuple): """A submitted-but-not-yet-drained transfer. @@ -325,7 +341,9 @@ class NixlTransport(DataTransport): transfer_id, exc, ) - continue + if not _is_remote_disconnect(exc): + continue + state = "ERR" if state == "DONE": if done_ids is None: done_ids = [] @@ -357,10 +375,14 @@ class NixlTransport(DataTransport): ) -> list[int]: """Cancel inflight transfers by their IDs. - See ``DataTransport.cancel`` for the contract. In "wait" mode, - transfers whose ``release_xfer_handle`` raises (NIXL could not - complete the abort because the backend is still draining) stay - in ``self._inflight`` so a later ``poll()`` will observe them. + See ``DataTransport.cancel`` for the contract. In "wait" mode, a + transfer is released only once NIXL reports it terminal. Releasing + a PROC/PEND handle does not stop the WRITE: the UCX backend frees + the request without cancelling RMA, and the bytes keep landing. So + PROC/PEND transfers, and those whose state check or + ``release_xfer_handle`` raises, stay in ``self._inflight`` for a + later ``poll()`` to observe. A state check that fails because the + peer disconnected counts as terminal. """ if mode == "immediate": handles = [ @@ -376,6 +398,27 @@ class NixlTransport(DataTransport): entry = self._inflight.get(tid) if entry is None: continue + try: + state = self._agent.check_xfer_state(entry.handle) + except Exception as exc: + logger.debug( + "NixlTransport %s: cancel state check failed for " + "transfer_id=%d: %s", + self._agent_name, + tid, + exc, + ) + if not _is_remote_disconnect(exc): + still_inflight.append(tid) + continue + # The handle may no longer be releasable; it is terminal + # either way, so drop it instead of retrying forever. + self._release_handles([entry.handle]) + del self._inflight[tid] + continue + if state in ("PROC", "PEND"): + still_inflight.append(tid) + continue try: self._agent.release_xfer_handle(entry.handle) except Exception as exc: diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index b7268eb420..a4b203548e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -1018,11 +1018,12 @@ class P2PSecondaryTierManager(SecondaryTierManager): def _drain_inflight_for_shutdown(self) -> None: """Best-effort drain of inflight transfers before closing _data. - Mirrors session._drain_abort but as a single bounded loop. Collects - inflight transfer_ids from each session, repeatedly calls + Collects inflight transfer_ids from each session, repeatedly calls _data.cancel(..., mode="wait") and _data.poll() so handles can - surface as done/failed, and falls back to mode="immediate" once - _SHUTDOWN_DRAIN_TIMEOUT_S elapses so a wedged peer can't hang us. + surface as done/failed. Unlike session._drain_abort, which never + force-releases, it falls back to mode="immediate" once + _SHUTDOWN_DRAIN_TIMEOUT_S elapses: the engine is exiting, and a + wedged peer must not hang shutdown. """ ids = [tid for s in self._sessions.values() for tid in s._server._inflight] if not ids: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 2e2c2b583f..fd43e31ec5 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -39,6 +39,10 @@ logger = init_logger(__name__) _STORE_TIMEOUT_S = 30.0 _CANCEL_DRAIN_TIMEOUT_S = 10.0 +# Minimum spacing between wait-mode cancel attempts for one round. Each +# attempt calls into NIXL on the scheduler thread, so an unconfirmed +# cancellation is retried at this rate instead of on every poll tick. +_CANCEL_RETRY_INTERVAL_S = 0.1 # Cap on time the server holds a HIT_PENDING / RETRY key from an inbound # LookupMsg before falling back to MISS. Long enough that an in-flight # primary write or a just-started promotion typically completes; short @@ -85,6 +89,13 @@ class _OutboundRequestState: remaining: int = 0 # blocks that need to be transferred to client finishing: bool = False # Signal finish request ASAP inflight: int = 0 # transfers submitted for this round, not yet polled + # The round was cancelled (a store job timed out, or the peer aborted + # it) while transfers may still be writing to the peer. Its jobs fail, + # but a job keeps its pin until every transfer carrying it is terminal: + # NIXL cannot stop a WRITE that is already in flight. + cancelled: bool = False + # Last wait-mode cancel attempt for this round's transfers (monotonic). + last_cancel_at: float | None = None # Job IDs that submit_store'd blocks for this round and have not # yet emitted a StoreResult. The terminal-finalize helper drains # this set; poll-done and poll-failed discard entries as their @@ -272,6 +283,9 @@ class ServerRole: # Parked aborts awaiting drain, keyed by (kv_request_id, round) # with the abort start time. self._pending_aborts: dict[tuple[str, int], float] = {} + # Parked aborts already warned about for exceeding + # _CANCEL_DRAIN_TIMEOUT_S. + self._slow_aborts: set[tuple[str, int]] = set() # ------------------------------------------------------------------ # State helpers @@ -322,6 +336,10 @@ class ServerRole: rnd = st.outbound.get(round_seq) if rnd is None: rnd = st.outbound[round_seq] = _OutboundRequestState() + if rnd.cancelled: + self._store_jobs.pop(job_id) + self._pending_store_results.append(StoreResult(job_id, success=False)) + return if from_lookup: rnd.lookup_supplied = True result = rnd.add_stored_blocks(keys, block_ids, job_id) @@ -765,7 +783,9 @@ class ServerRole: tid, ) continue - results.extend(self._settle_xfer_jobs(xfer, success=True)) + results.extend( + self._settle_xfer_jobs(xfer, success=not xfer.round.cancelled) + ) rnd = xfer.round st = self._requests.get(xfer.kv_request_id) if st is not None and st.outbound.get(xfer.round_key) is rnd: @@ -773,7 +793,10 @@ class ServerRole: assert rnd.remaining >= 0, ( f"remaining went negative for kv_request_id={xfer.kv_request_id}" ) - if rnd.remaining == 0: + if rnd.cancelled: + if rnd.inflight == 0: + self._finalize_cancelled(xfer.kv_request_id, xfer.round_key) + elif rnd.remaining == 0: self._finalize_outbound( xfer.kv_request_id, xfer.round_key, success=True ) @@ -783,7 +806,6 @@ class ServerRole: ) self._maybe_prune(xfer.kv_request_id) - failed_rounds: list[tuple[str, _OutboundRequestState]] | None = None for tid in poll_result.failed: xfer = self._inflight_pop(tid) if xfer is None: @@ -796,38 +818,21 @@ class ServerRole: tid, ) continue - results.extend(self._settle_xfer_jobs(xfer, success=False)) rnd = xfer.round st = self._requests.get(xfer.kv_request_id) - if st is not None and st.outbound.get(xfer.round_key) is rnd: - del st.outbound[xfer.round_key] - if failed_rounds is None: - failed_rounds = [] - failed_rounds.append((xfer.kv_request_id, rnd)) - self._send( - { - TYPE_KEY: TransferDoneMsg.TYPE, - TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id, - TransferDoneMsg.SUCCESS: False, - TransferDoneMsg.ROUND_SEQ: xfer.round_key, - } - ) + in_round = st is not None and st.outbound.get(xfer.round_key) is rnd + if in_round and not rnd.cancelled: + # One failed transfer fails the round, but its other + # transfers may still be writing. Cancel them like a + # timed-out round: their pins and the TransferDoneMsg wait + # until they are terminal. + rnd.cancelled = True + results.extend(self._fail_unwritten_round_jobs(rnd)) + results.extend(self._settle_xfer_jobs(xfer, success=False)) + if in_round and rnd.inflight == 0: + self._finalize_cancelled(xfer.kv_request_id, xfer.round_key) self._maybe_prune(xfer.kv_request_id) - # Cancel each failed round's other inflight and fail its - # remaining store jobs — nothing else will settle them. - if failed_rounds: - for kv_request_id, rnd in failed_rounds: - ids_to_cancel = [ - tid for tid, x in self._inflight.items() if x.round is rnd - ] - for tid in ids_to_cancel: - self._inflight_pop(tid) - if ids_to_cancel: - self._transport.cancel(ids_to_cancel) - results.extend(self._fail_round_jobs(rnd)) - self._maybe_prune(kv_request_id) - return results def collect_idle_timeouts(self) -> list[StoreResult]: @@ -870,6 +875,7 @@ class ServerRole: self._requests.clear() self._serve_pending.clear() self._pending_aborts.clear() + self._slow_aborts.clear() self._finished_lookup_ctxs.clear() return failed_stores, failed_serves @@ -918,7 +924,10 @@ class ServerRole: result. """ results: list[StoreResult] = [] + still_pinned = self._read_job_ids() if xfer.round.cancelled else set() for job_id in xfer.job_ids: + if job_id in still_pinned: + continue if self._store_jobs.pop(job_id, None) is None: continue results.append(StoreResult(job_id=job_id, success=success)) @@ -939,6 +948,48 @@ class ServerRole: rnd.pending_job_ids.clear() return results + def _finalize_cancelled(self, kv_request_id: str, round_key: int) -> None: + """Finalize a cancelled round once its last transfer is terminal. + + An aborted round gets no TransferDoneMsg: its parked abort answers + the peer with AbortAckMsg on the next drain. + """ + self._finalize_outbound( + kv_request_id, + round_key, + success=False, + send_done=(kv_request_id, round_key) not in self._pending_aborts, + ) + + def _read_job_ids(self, rnd: _OutboundRequestState | None = None) -> set[int]: + """Jobs whose blocks an inflight transfer (of ``rnd``, if given) + still reads; their pins must be held.""" + return { + jid + for xfer in self._inflight.values() + if rnd is None or xfer.round is rnd + for jid in xfer.job_ids + } + + def _fail_unwritten_round_jobs( + self, rnd: _OutboundRequestState + ) -> list[StoreResult]: + """Fail a cancelled round's jobs that no inflight transfer carries. + + Jobs that an inflight transfer carries keep their pins; they settle + when that transfer turns terminal (``_settle_xfer_jobs``). + """ + written = self._read_job_ids(rnd) + results: list[StoreResult] = [] + for job_id in list(rnd.pending_job_ids): + if job_id in written: + continue + rnd.pending_job_ids.discard(job_id) + if self._store_jobs.pop(job_id, None) is None: + continue + results.append(StoreResult(job_id=job_id, success=False)) + return results + def _finalize_outbound( self, kv_request_id: str, @@ -953,7 +1004,12 @@ class ServerRole: fetch). Other rounds of the id are untouched. """ st = self._requests[kv_request_id] - req = st.outbound.pop(round_key) + req = st.outbound[round_key] + if req.cancelled and req.inflight: + return + del st.outbound[round_key] + if req.cancelled: + success = False if success is None: success = req.demand_received and req.remaining == 0 settled = self._fail_round_jobs(req) if not success else None @@ -992,41 +1048,56 @@ class ServerRole: def _drain_abort(self, kv_request_id: str, round_seq: int) -> None: """One drain attempt for a pending abort. - Detaches the aborted round, then asks the transport to cancel its - inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` once - nothing remains inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` - falls back to ``mode="immediate"`` and acks anyway. + Cancels the aborted round, then asks the transport to cancel its + inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` only + once none of them can still write to the peer: the ack lets the + peer reuse the destination blocks. A transfer that NIXL does not + report terminal is never force-released, however long the drain + takes: neither releasing its handle nor removing the remote agent + stops a WRITE that is already in flight. Past + ``_CANCEL_DRAIN_TIMEOUT_S`` the drain only logs a warning. """ st = self._requests[kv_request_id] - rnd = st.outbound.pop(round_seq, None) - if rnd is not None: - # Its transfers are being cancelled; fail its jobs now - # instead of leaking them to the store timeout. - self._pending_store_results.extend(self._fail_round_jobs(rnd)) + rnd = st.outbound.get(round_seq) + if rnd is not None and not rnd.cancelled: + # Fail the round's jobs now instead of leaking them to the + # store timeout, but keep the pins of those still being written. + rnd.cancelled = True + self._pending_store_results.extend(self._fail_unwritten_round_jobs(rnd)) ids = [ tid for tid, x in self._inflight.items() if x.kv_request_id == kv_request_id and x.round_key == round_seq ] if not ids: + if rnd is not None and st.outbound.get(round_seq) is rnd: + self._finalize_outbound( + kv_request_id, round_seq, success=False, send_done=False + ) self._finalize_abort(kv_request_id, round_seq) return - started_at = self._pending_aborts[(kv_request_id, round_seq)] - if time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S: - for tid in ids: - self._inflight_pop(tid) - self._transport.cancel(ids, mode="immediate") + now = time.monotonic() + key = (kv_request_id, round_seq) + if ( + now - self._pending_aborts[key] >= _CANCEL_DRAIN_TIMEOUT_S + and key not in self._slow_aborts + ): + self._slow_aborts.add(key) logger.warning( - "P2PSession %s: cancel drain timed out for kv_request_id=%s," - " force-canceled %d transfers", + "P2PSession %s: abort of kv_request_id=%s round=%s still " + "waiting after %.0fs for %d transfers to stop; holding their " + "pins and the AbortAck until NIXL reports them terminal", self._peer_id, kv_request_id, + round_seq, + _CANCEL_DRAIN_TIMEOUT_S, len(ids), ) - self._finalize_abort(kv_request_id, round_seq) - return + # All of the aborted round's transfers share one round state. + if not self._cancel_due(self._inflight[ids[0]].round, now): + return still = self._transport.cancel(ids, mode="wait") # Tids the transport successfully released are gone from its # _inflight; mirror that in session bookkeeping so they don't @@ -1037,10 +1108,30 @@ class ServerRole: if tid not in still_set: self._inflight_pop(tid) if not still: + if rnd is not None and st.outbound.get(round_seq) is rnd: + self._finalize_outbound( + kv_request_id, round_seq, success=False, send_done=False + ) self._finalize_abort(kv_request_id, round_seq) + @staticmethod + def _cancel_due(rnd: _OutboundRequestState, now: float) -> bool: + """Rate-limit wait-mode cancels of ``rnd``'s transfers. + + Shared by the store-timeout sweep and the abort drain, so a round + gets at most one attempt per ``_CANCEL_RETRY_INTERVAL_S``. + """ + if ( + rnd.last_cancel_at is not None + and now - rnd.last_cancel_at < _CANCEL_RETRY_INTERVAL_S + ): + return False + rnd.last_cancel_at = now + return True + def _finalize_abort(self, kv_request_id: str, round_seq: int) -> None: self._pending_aborts.pop((kv_request_id, round_seq), None) + self._slow_aborts.discard((kv_request_id, round_seq)) self._send( { TYPE_KEY: AbortAckMsg.TYPE, @@ -1121,7 +1212,8 @@ class ServerRole: def _timeout_pending_store_jobs(self) -> list[StoreResult]: if not self._store_jobs: return [] - deadline = time.monotonic() - _STORE_TIMEOUT_S + now = time.monotonic() + deadline = now - _STORE_TIMEOUT_S timed_out: list[int] | None = None for jid, submitted_at in self._store_jobs.items(): if submitted_at <= deadline: @@ -1129,10 +1221,44 @@ class ServerRole: timed_out = [] timed_out.append(jid) if timed_out is None: - return [] + if not any(xfer.round.cancelled for xfer in self._inflight.values()): + return [] + timed_out = [] results: list[StoreResult] = [] + expired = set(timed_out) + rounds = { + (xfer.kv_request_id, xfer.round_key): xfer.round + for xfer in self._inflight.values() + if xfer.round.cancelled or expired.intersection(xfer.job_ids) + } + for (kv_request_id, round_key), rnd in rounds.items(): + rnd.cancelled = True + if not self._cancel_due(rnd, now): + continue + tids = [tid for tid, xfer in self._inflight.items() if xfer.round is rnd] + still = set(self._transport.cancel(tids, mode="wait")) + for tid in tids: + if tid in still: + continue + xfer = self._inflight_pop(tid) + assert xfer is not None + results.extend(self._settle_xfer_jobs(xfer, success=False)) + st = self._requests.get(kv_request_id) + if ( + rnd.inflight == 0 + and st is not None + and st.outbound.get(round_key) is rnd + ): + self._finalize_cancelled(kv_request_id, round_key) + + # Cancellation may still be draining. Keep every affected pin until + # the transport confirms that no transfer reads its primary slot. + inflight_job_ids = self._read_job_ids() for jid in timed_out: - del self._store_jobs[jid] + if jid in inflight_job_ids: + continue + if self._store_jobs.pop(jid, None) is None: + continue results.append(StoreResult(job_id=jid, success=False)) logger.warning("P2PSession %s: store job %d timed out", self._peer_id, jid) # A failed store job's primary slots are unpinned by the engine, so its