[KVConnector][P2P] Register peers off the scheduler thread (#57139)

Signed-off-by: Liran Schour <lirans@il.ibm.com>
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 <oro@il.ibm.com>

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()

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"

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()

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.

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()])

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,