From a2fb4827324bf023075799ef3d70d347834b8137 Mon Sep 17 00:00:00 2001 From: liranschour Date: Wed, 14 Oct 2026 08:21:57 +0300 Subject: [PATCH] [KVConnector][P2P] Register peers off the scheduler thread (#57139) Signed-off-by: Liran Schour Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../tiering/p2p/test_data_transport.py | 55 ++++++++ .../v1/kv_offload/tiering/p2p/test_manager.py | 14 ++ .../kv_offload/tiering/p2p/test_sessions.py | 102 +++++++++++++- vllm/v1/kv_offload/tiering/p2p/data/base.py | 48 +++++++ vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 128 ++++++++++++++++-- .../kv_offload/tiering/p2p/session/session.py | 90 ++++++++++-- 6 files changed, 416 insertions(+), 21 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 e40f16ad73..60830a34bb 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -5,9 +5,11 @@ from __future__ import annotations import ctypes +import threading from unittest.mock import MagicMock, patch import numpy as np +import pytest from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport @@ -107,6 +109,59 @@ class TestNixlTransportWithMockedAgent: transport = self._make_transport() assert transport.get_agent_metadata() == b"test-metadata" + def test_add_remote_peer_async_registers_off_the_caller_thread(self): + """The O(num_blocks) registration must not run on the caller thread. + + This is the whole point of the async path: the scheduler thread + hands the work off rather than paying prep_xfer_dlist inline. + """ + transport = self._make_transport() + caller = threading.get_ident() + worker_threads: list[int] = [] + + def record_thread(*args, **kwargs): + worker_threads.append(threading.get_ident()) + return MagicMock(name="remote_dlist") + + transport._agent.prep_xfer_dlist.side_effect = record_thread + + future = transport.add_remote_peer_async("peer:1", b"meta", 0x1000, 8, 1024) + future.result(timeout=5) + + assert worker_threads and caller not in worker_threads + assert transport.write_blocks("peer:1", [0], [1]) is not None + + def test_registration_landing_after_remove_is_discarded(self): + """A late registration must not revive a peer that was reaped. + + The session is torn down and remove_remote_peer runs while the + worker is still building descriptors; publishing afterwards would + leave a registered peer with no session and leak its handles. + """ + transport = self._make_transport() + building = threading.Event() + release = threading.Event() + + def blocking_prep(*args, **kwargs): + building.set() + release.wait(timeout=5) + return MagicMock(name="remote_dlist") + + transport._agent.prep_xfer_dlist.side_effect = blocking_prep + + future = transport.add_remote_peer_async("peer:1", b"meta", 0x1000, 8, 1024) + assert building.wait(timeout=5) + transport.remove_remote_peer("peer:1") + release.set() + + with pytest.raises(RuntimeError, match="superseded"): + future.result(timeout=5) + assert "peer:1" not in transport._remote_dlists + assert transport.write_blocks("peer:1", [0], [1]) is None + # The handles the worker built were released rather than leaked. + assert transport._agent.release_dlist_handle.called + assert transport._agent.remove_remote_agent.called + def test_write_blocks_returns_none_for_unknown_peer(self): """write_blocks returns None if peer not registered.""" transport = self._make_transport() diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 930158c7f6..0ca2735005 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -10,6 +10,7 @@ from __future__ import annotations import time import uuid +from concurrent.futures import Future from types import SimpleNamespace import numpy as np @@ -1111,6 +1112,14 @@ class _FakeData: "block_len": block_len, } + def add_remote_peer_async( + self, peer_id, agent_metadata, base_addr, num_blocks, block_len + ): + future: Future[None] = Future() + self.add_remote_peer(peer_id, agent_metadata, base_addr, num_blocks, block_len) + future.set_result(None) + return future + def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) @@ -1368,6 +1377,11 @@ class TestAcceptNewPeers: def add_remote_peer(self, *args, **kwargs): pass + def add_remote_peer_async(self, *args, **kwargs): + future: Future[None] = Future() + future.set_result(None) + return future + mgr = _make_manager() mgr._data = FakeData() peer_id = "10.0.0.1:8000" diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 300cbacd66..0d756cbf42 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -13,7 +13,8 @@ completes its own load. from __future__ import annotations import time -from collections.abc import Sequence +from collections.abc import Callable, Sequence +from concurrent.futures import Future import numpy as np import pytest @@ -74,7 +75,13 @@ class FakeDataTransport: num_blocks: int = 16, block_len: int = 4096, config_fingerprint: str = "", + defer_registration: bool = False, ) -> None: + # When True, add_remote_peer_async parks the registration until + # complete_registrations() resolves it, standing in for the real + # transport's worker thread. + self.defer_registration = defer_registration + self._deferred: list[tuple[Future[None], Callable[[], None]]] = [] self._base_addr = base_addr self._num_blocks = num_blocks self._block_len = block_len @@ -116,6 +123,38 @@ class FakeDataTransport: "block_len": block_len, } + def add_remote_peer_async( + self, peer_id, agent_metadata, base_addr, num_blocks, block_len + ) -> Future[None]: + future: Future[None] = Future() + + def apply() -> None: + self.add_remote_peer( + peer_id, agent_metadata, base_addr, num_blocks, block_len + ) + + if self.defer_registration: + self._deferred.append((future, apply)) + return future + apply() + future.set_result(None) + return future + + def complete_registrations(self, exc: Exception | None = None) -> int: + """Resolve parked registrations the way the worker thread would. + + Returns the number resolved so callers can assert one was actually + in flight rather than silently testing nothing. + """ + deferred, self._deferred = self._deferred, [] + for future, apply in deferred: + if exc is not None: + future.set_exception(exc) + else: + apply() + future.set_result(None) + return len(deferred) + def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) @@ -390,6 +429,67 @@ class TestConnectHandshake: ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) assert ack[ConnectAckMsg.PEER_ID] == "local:9000" + def test_ack_withheld_until_registration_completes(self): + """ConnectAck must not precede registration. + + The peer treats ConnectAck as permission to send, so acking while + registration is still in flight would let a FetchMsg arrive for + blocks this side cannot yet write to. + """ + transport = FakeDataTransport(defer_registration=True) + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + + assert "peer:8000" not in transport._remote_peers + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + assert transport.complete_registrations() == 1 + session.poll() + + assert "peer:8000" in transport._remote_peers + ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) + assert ack[ConnectAckMsg.PEER_ID] == "local:9000" + + def test_pending_registration_is_pending_work(self): + """An in-flight registration keeps the engine ticking. + + Nothing else drives poll(), so if the session reported idle the + ack would never be sent. + """ + transport = FakeDataTransport(defer_registration=True) + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + assert session.has_pending_work + + transport.complete_registrations() + session.poll() + assert not session.has_pending_work + + def test_failed_registration_rejects_peer(self): + """A registration that fails kills the connection instead of acking.""" + transport = FakeDataTransport(defer_registration=True) + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + + transport.complete_registrations(exc=RuntimeError("nixl exploded")) + session.poll() + + assert not conn.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_duplicate_connect_during_registration_is_rejected(self): + """A second ConnectMsg mid-registration is a protocol violation.""" + transport = FakeDataTransport(defer_registration=True) + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + conn.enqueue(_peer_connect_msg()) + session.poll() + assert not conn.alive + def test_connect_ack_makes_session_ready(self): """Session.ready becomes True after ConnectAckMsg.""" session, conn, _ = _make_session() diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index 62304b7544..574ba8955d 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -76,6 +76,7 @@ import hashlib import json from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence +from concurrent.futures import Future from typing import Literal, NamedTuple CancelMode = Literal["immediate", "wait"] @@ -179,6 +180,53 @@ class DataTransport(ABC): """ ... + def add_remote_peer_async( + self, + peer_id: str, + agent_metadata: bytes, + base_addr: int, + num_blocks: int, + block_len: int, + ) -> Future[None]: + """Register a remote peer without blocking the caller. + + Registration is O(num_blocks) and can take seconds for a large + secondary tier, so the session starts it here and sends + ``ConnectAckMsg`` only once the returned future resolves. The + acknowledgement is therefore the peer's guarantee that this side + is ready to serve, and no ``write_blocks`` can arrive before it. + + The default implementation runs ``add_remote_peer`` inline and + returns an already-resolved future; transports whose registration + is cheap need not override it. + + Args: + peer_id: Unique identifier for the remote peer. + agent_metadata: Opaque bytes from the peer's get_agent_metadata(). + base_addr: Base address of the peer's block memory region. + num_blocks: Number of blocks in the peer's region. + block_len: Size of each block (must match local block_len). + + Returns: + A future resolving to None on success. A failed future means + the peer must be rejected rather than acknowledged. + + """ + future: Future[None] = Future() + try: + self.add_remote_peer( + peer_id, + agent_metadata=agent_metadata, + base_addr=base_addr, + num_blocks=num_blocks, + block_len=block_len, + ) + except Exception as exc: + future.set_exception(exc) + else: + future.set_result(None) + return future + @abstractmethod def remove_remote_peer(self, peer_id: str) -> None: """Unregister a remote peer and release associated resources. diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index e8fb66dd6e..c44069ceec 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -5,7 +5,9 @@ from __future__ import annotations import itertools +import threading from collections.abc import Iterable +from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, NamedTuple import numpy as np @@ -66,6 +68,23 @@ class NixlTransport(DataTransport): # transfer_id → _Inflight(peer_id, handle). self._inflight: dict[int, _Inflight] = {} self._next_id = itertools.count() + # Registration is O(num_blocks) NIXL work — add_remote_agent plus a + # prep_xfer_dlist over every block of the peer's region — so it runs + # off the scheduler thread, mirroring the NIXL connector's + # _handshake_initiation_executor. One worker only: NIXL is not + # guaranteed to be thread-safe, so registrations stay serialized + # against each other. + self._reg_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="vllm-p2p-nixl-reg" + ) + # peer_id → registration generation, bumped on every add and remove. + # A worker that finishes building handles under a stale generation + # releases them instead of publishing, so a late completion cannot + # revive a peer that was reaped or superseded meanwhile. + self._peer_gen: dict[str, int] = {} + # Guards _peer_gen together with _peer_nixl_names/_remote_dlists so + # the generation check and the publish are one atomic step. + self._peer_lock = threading.Lock() self._init(view) @@ -126,23 +145,113 @@ class NixlTransport(DataTransport): num_blocks: int, block_len: int, ) -> None: + gen = self._begin_registration(peer_id) + nixl_name, remote_dlist = self._build_peer_handles( + agent_metadata, base_addr, num_blocks, block_len + ) + self._publish_peer(peer_id, gen, nixl_name, remote_dlist) + + def add_remote_peer_async( + self, + peer_id: str, + agent_metadata: bytes, + base_addr: int, + num_blocks: int, + block_len: int, + ) -> Future[None]: + gen = self._begin_registration(peer_id) + + def register() -> None: + nixl_name, remote_dlist = self._build_peer_handles( + agent_metadata, base_addr, num_blocks, block_len + ) + if not self._publish_peer(peer_id, gen, nixl_name, remote_dlist): + raise RuntimeError( + f"registration for peer {peer_id} was superseded " + "(peer removed or re-registered while registering)" + ) + + return self._reg_executor.submit(register) + + def remove_remote_peer(self, peer_id: str) -> None: + with self._peer_lock: + # Invalidate any registration still being built for this peer so + # its worker releases the handles instead of publishing them. + self._peer_gen[peer_id] = self._peer_gen.get(peer_id, 0) + 1 + nixl_name = self._peer_nixl_names.pop(peer_id, None) + dlist = self._remote_dlists.pop(peer_id, None) + if self._agent is not None: + if dlist is not None: + self._agent.release_dlist_handle(dlist) + if nixl_name: + self._agent.remove_remote_agent(nixl_name) + + def _begin_registration(self, peer_id: str) -> int: + """Claim the next registration generation for *peer_id*.""" + with self._peer_lock: + gen = self._peer_gen.get(peer_id, 0) + 1 + self._peer_gen[peer_id] = gen + return gen + + def _build_peer_handles( + self, + agent_metadata: bytes, + base_addr: int, + num_blocks: int, + block_len: int, + ) -> tuple[str, Any]: + """The expensive half: agent wireup plus one descriptor per block. + + Touches no shared state, so it runs unlocked — holding the peer + lock across it would stall the scheduler thread for exactly as + long as the registration it is meant to move off that thread. + """ nixl_name = self._agent.add_remote_agent(agent_metadata) block_descs = [ (base_addr + i * block_len, block_len, 0) for i in range(num_blocks) ] xfer_dlist = self._agent.get_xfer_descs(block_descs, mem_type="DRAM") remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist) - self._peer_nixl_names[peer_id] = nixl_name - self._remote_dlists[peer_id] = remote_dlist + return nixl_name, remote_dlist - def remove_remote_peer(self, peer_id: str) -> None: - nixl_name = self._peer_nixl_names.pop(peer_id, None) - dlist = self._remote_dlists.pop(peer_id, None) + def _publish_peer( + self, + peer_id: str, + gen: int, + nixl_name: str, + remote_dlist: Any, + ) -> bool: + """Install freshly built handles, unless *gen* is stale. + + Returns False when the peer was removed or re-registered while the + handles were being built; the handles are released here since + nothing else ever saw them. + """ + with self._peer_lock: + fresh = self._peer_gen.get(peer_id) == gen + if fresh: + self._peer_nixl_names[peer_id] = nixl_name + self._remote_dlists[peer_id] = remote_dlist + if fresh: + return True + logger.info( + "NixlTransport %s: discarding superseded registration for peer=%s", + self._agent_name, + peer_id, + ) if self._agent is not None: - if dlist is not None: - self._agent.release_dlist_handle(dlist) - if nixl_name: + try: + self._agent.release_dlist_handle(remote_dlist) self._agent.remove_remote_agent(nixl_name) + except Exception as exc: + logger.warning( + "NixlTransport %s: releasing superseded registration " + "for peer=%s failed: %s", + self._agent_name, + peer_id, + exc, + ) + return False # ------------------------------------------------------------------ # Transfer submission and polling @@ -286,6 +395,9 @@ class NixlTransport(DataTransport): # ------------------------------------------------------------------ def close(self) -> None: + # Stop accepting registrations and let any in-flight one finish + # before the agent goes away underneath it. + self._reg_executor.shutdown(wait=True) if self._agent is None: return self._release_handles([entry.handle for entry in self._inflight.values()]) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index ef490304c9..6d3d990034 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -11,8 +11,11 @@ dispatch — each parsed message is forwarded to the corresponding role. Wire protocol is unchanged. Both sides advertise their NIXL metadata via ConnectMsg when their session is connected; the peer's ConnectMsg -triggers transport.add_remote_peer; ConnectAckMsg confirms the peer -received our ConnectMsg, after which queued outgoing messages are flushed. +starts transport.add_remote_peer_async off the scheduler thread, and +ConnectAckMsg is sent only once that registration lands. ConnectAckMsg +confirms the peer received our ConnectMsg, after which queued outgoing +messages are flushed — so a peer that may send is a peer we have already +registered. """ from __future__ import annotations @@ -20,6 +23,7 @@ from __future__ import annotations import contextlib import time from collections.abc import Sequence +from concurrent.futures import Future from typing import TYPE_CHECKING, NamedTuple import vllm.envs as envs @@ -100,9 +104,11 @@ class P2PSession: the decoder connects). - Constructor with conn != None ⇒ connected. Sends our own ConnectMsg immediately; the peer's ConnectMsg arrives in poll() and is - dispatched to _on_connect (which calls transport.add_remote_peer - and replies with ConnectAckMsg). Outgoing sends are queued until - ConnectAckMsg confirms our metadata reached the peer. + dispatched to _on_connect, which starts + transport.add_remote_peer_async. A later poll() sees the + registration land and replies with ConnectAckMsg; a failed + registration marks the connection dead instead. Outgoing sends are + queued until ConnectAckMsg confirms our metadata reached the peer. - attach_connection(conn) on a pending session ⇒ same as above, starting from pending. """ @@ -140,6 +146,13 @@ class P2PSession: # submit_store batches parked before the binding existed. self._new_fetch_ids: list[str] = [] + # In-flight peer registration started by _on_connect. Registration + # is O(num_blocks) and runs off this thread, so ConnectAckMsg is + # held back until the future resolves — the ack is what tells the + # peer it may start sending, and nothing may be served before the + # peer's blocks are registered here. + self._pending_registration: Future[None] | None = None + self._client = ClientRole(peer_id=peer_id, send=self._send) self._server = ServerRole( peer_id=peer_id, @@ -172,7 +185,11 @@ class P2PSession: @property def has_pending_work(self) -> bool: """True while inbound loads or outbound transfers are outstanding.""" - return self._client.has_active_loads or self._server.has_inflight_transfers + return ( + self._pending_registration is not None + or self._client.has_active_loads + or self._server.has_inflight_transfers + ) @property def last_activity_at(self) -> float: @@ -278,9 +295,21 @@ class P2PSession: for msg in messages: self._on_message(msg) + # After dispatch: a ConnectMsg handled just above may already have a + # resolved registration (transports that register inline ack in the + # same tick), and a registration started on an earlier tick may have + # landed while the engine was busy elsewhere. This runs before the + # handshake deadline below so a registration that just resolved gets + # its ConnectAckMsg out rather than being judged as timed out, and so + # a failed registration's mark_dead() short-circuits that check. + self._poll_registration() + # Handshake deadline: a peer that never completes the handshake # is stuck in connected-but-not-ready state. Mark it dead so the # manager reaps it instead of accumulating sockets indefinitely. + # Note this also covers the peer registering *our* region before it + # acks, so the deadline must exceed remote registration time; the + # 30s default leaves ample room for the descriptor build. if ( not self._send_ready and self._conn.alive @@ -320,6 +349,14 @@ class P2PSession: ``parent.on_request_finished`` (the manager flushes these on its next ``serve_external_requests``). """ + # Abandon any in-flight registration: nothing will ack this peer + # now. The manager's remove_remote_peer bumps the transport's + # generation for this peer, so a worker already past cancellation + # releases its handles instead of publishing them. + if self._pending_registration is not None: + self._pending_registration.cancel() + self._pending_registration = None + client_result = self._client.close() failed_stores, failed_serves = self._server.close() @@ -470,10 +507,11 @@ class P2PSession: # Validation failures here mean an incompatible or malicious peer. # Mark the connection dead so the manager reaps the session; # don't call add_remote_peer or send connect_ack. - if self._send_ready: - # We've already received connect_ack, so the handshake is - # complete. A second connect from the peer is a protocol - # violation — re-registering would corrupt transport state. + if self._send_ready or self._pending_registration is not None: + # Either the handshake completed (we hold the peer's ack) or a + # registration for this peer is already in flight. A second + # connect either way is a protocol violation — re-registering + # would corrupt transport state. self._protocol_error("duplicate connect after handshake") return try: @@ -498,7 +536,7 @@ class P2PSession: f"local={self._local_hash_seed!r}. Ensure PYTHONHASHSEED " "(if set) matches on all P2P peers." ) - self._transport.add_remote_peer( + registration = self._transport.add_remote_peer_async( self.peer_id, agent_metadata=msg[ConnectMsg.AGENT_METADATA], base_addr=msg[ConnectMsg.BASE_ADDR], @@ -511,7 +549,35 @@ class P2PSession: self._conn.mark_dead() return - if self._conn is not None: + # ConnectAckMsg is deferred to _poll_registration once the + # registration lands; sending it now would let the peer fetch + # against blocks this side has not registered yet. + self._pending_registration = registration + + def _poll_registration(self) -> None: + """Ack or reject the peer once its registration resolves. + + Called from poll() after message dispatch. A successful + registration releases ConnectAckMsg — the peer treats that as + permission to send, so it must not go out earlier. A failed one + marks the connection dead so the manager reaps the session. + """ + registration = self._pending_registration + if registration is None or not registration.done(): + return + self._pending_registration = None + try: + registration.result() + except Exception as exc: + logger.error( + "P2PSession %s: peer registration failed, rejecting peer: %s", + self.peer_id, + exc, + ) + if self._conn is not None: + self._conn.mark_dead() + return + if self._conn is not None and self._conn.alive: self._conn.send( { TYPE_KEY: ConnectAckMsg.TYPE,