diff --git a/tests/v1/engine/test_connector_poller.py b/tests/v1/engine/test_connector_poller.py new file mode 100644 index 000000000000..d3947e4c19f2 --- /dev/null +++ b/tests/v1/engine/test_connector_poller.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ConnectorPoller: services a connector only while the engine waits.""" + +import threading +import time +from concurrent.futures import Future + +import pytest + +from vllm.v1.engine.core import ConnectorPoller + + +class _CountingConnector: + def __init__(self, fail: bool = False) -> None: + self.calls = 0 + self.fail = fail + self.thread_names: set[str] = set() + + def poll_pending_work(self) -> None: + self.calls += 1 + self.thread_names.add(threading.current_thread().name) + if self.fail: + raise RuntimeError("sweep failed") + + +def _resolve_after(future: Future, delay_s: float, value: object = "out") -> None: + def run(): + time.sleep(delay_s) + future.set_result(value) + + threading.Thread(target=run, daemon=True).start() + + +def test_poller_sweeps_only_while_waiting(): + connector = _CountingConnector() + poller = ConnectorPoller(connector, interval_s=0.001) + try: + time.sleep(0.02) + assert connector.calls == 0 + + future: Future = Future() + _resolve_after(future, 0.05) + assert poller.wait(future) == "out" + assert connector.calls > 0 + assert connector.thread_names == {"kv-connector-poller"} + + settled = connector.calls + time.sleep(0.02) + assert connector.calls == settled + finally: + poller.close() + + +def test_poller_reraises_sweep_error_on_engine_thread(): + connector = _CountingConnector(fail=True) + poller = ConnectorPoller(connector, interval_s=0.001) + try: + future: Future = Future() + _resolve_after(future, 0.03) + with pytest.raises(RuntimeError, match="sweep failed"): + poller.wait(future) + assert connector.calls == 1 + + # A later wait runs sweeps again; the stored error was consumed. + connector.fail = False + future = Future() + _resolve_after(future, 0.03) + assert poller.wait(future) == "out" + assert connector.calls > 1 + finally: + poller.close() + + +def test_poller_propagates_future_exception(): + connector = _CountingConnector() + poller = ConnectorPoller(connector, interval_s=0.001) + try: + future: Future = Future() + future.set_exception(ValueError("model failed")) + with pytest.raises(ValueError, match="model failed"): + poller.wait(future) + finally: + poller.close() 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 d3bacc8f326f..15aefb10f2c2 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -324,7 +324,9 @@ def test_non_ucx_backends_passes_backends_kwarg(self): ): NixlTransport("test:1", self._make_view(), backends=["MOONCAKE"]) - config_fn.assert_called_once_with(backends=["MOONCAKE"], capture_telemetry=True) + config_fn.assert_called_once_with( + backends=["MOONCAKE"], capture_telemetry=True, sync_mode=None + ) # num_threads must NOT be passed on the non-UCX branch. assert "num_threads" not in config_fn.call_args.kwargs @@ -340,7 +342,9 @@ def test_ucx_only_passes_num_threads(self): ): NixlTransport("test:1", self._make_view(), num_threads=8) - config_fn.assert_called_once_with(num_threads=8, capture_telemetry=True) + config_fn.assert_called_once_with( + num_threads=8, capture_telemetry=True, sync_mode=None + ) assert "backends" not in config_fn.call_args.kwargs def test_default_backends_is_ucx_only(self): @@ -356,4 +360,30 @@ def test_default_backends_is_ucx_only(self): NixlTransport("test:1", self._make_view()) # Default num_threads=4, no backends kwarg. - config_fn.assert_called_once_with(num_threads=4, capture_telemetry=True) + config_fn.assert_called_once_with( + num_threads=4, capture_telemetry=True, sync_mode=None + ) + + def test_agent_runs_in_strict_thread_sync_when_available(self): + """With a NIXL that exposes the sync enum, the agent serializes its + API so peer registration may run on the worker thread.""" + import sys + import types + + fake_module = types.ModuleType("fake_nixl_api") + fake_module.nixl_thread_sync_t = types.SimpleNamespace( + NIXL_THREAD_SYNC_STRICT="strict" + ) + config_fn = MagicMock(return_value=MagicMock(name="cfg")) + config_fn.__module__ = fake_module.__name__ + agent_cls = MagicMock() + with ( + patch.dict(sys.modules, {fake_module.__name__: fake_module}), + patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls), + patch( + "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn + ), + ): + NixlTransport("test:1", self._make_view()) + + assert config_fn.call_args.kwargs["sync_mode"] == "strict" diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index aceec1cd8fa1..8190a9f8b1d0 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -1073,6 +1073,16 @@ def add_remote_peer( "block_len": block_len, } + def add_remote_peer_async( + self, peer_id, agent_metadata, base_addr, num_blocks, block_len + ): + from concurrent.futures import Future + + self.add_remote_peer(peer_id, agent_metadata, base_addr, num_blocks, block_len) + future: Future[None] = Future() + future.set_result(None) + return future + def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 2e4415fbb75d..be83dd4c71c5 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -14,6 +14,7 @@ import time from collections.abc import Sequence +from concurrent.futures import Future import numpy as np import pytest @@ -86,6 +87,10 @@ def __init__( self._poll_failed: list[int] = [] self._cancel_still_inflight: set[int] = set() self._cancel_calls: list[tuple[list[int], str]] = [] + # When True, add_remote_peer_async parks the registration in + # pending_registrations for the test to complete or fail. + self.defer_registration = False + self.pending_registrations: dict[str, Future[None]] = {} @property def base_addr(self) -> int: @@ -116,6 +121,25 @@ def add_remote_peer( "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() + if self.defer_registration: + self.pending_registrations[peer_id] = future + future.add_done_callback( + lambda f: ( + f.exception() is None + and self.add_remote_peer( + peer_id, agent_metadata, base_addr, num_blocks, block_len + ) + ) + ) + return 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) @@ -390,6 +414,41 @@ def test_peer_connect_triggers_add_remote_and_ack(self): ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) assert ack[ConnectAckMsg.PEER_ID] == "local:9000" + def test_connect_ack_waits_for_peer_registration(self): + """ConnectAck is deferred until the transport has registered the peer.""" + transport = FakeDataTransport() + transport.defer_registration = True + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + assert "peer:8000" not in transport._remote_peers + assert session.has_pending_work + session.poll() + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + transport.pending_registrations["peer:8000"].set_result(None) + 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" + assert not session.has_pending_work + + def test_failed_peer_registration_rejects_peer(self): + """A registration error rejects the peer like a validation failure.""" + transport = FakeDataTransport() + transport.defer_registration = True + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + transport.pending_registrations["peer:8000"].set_exception( + RuntimeError("agent metadata rejected") + ) + session.poll() + assert not session.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + assert not session.has_pending_work + def test_connect_ack_makes_session_ready(self): """Session.ready becomes True after ConnectAckMsg.""" session, conn, _ = _make_session() diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 771153ec4734..5fb78ba29422 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -148,6 +148,46 @@ def get_stats(self) -> OffloadingConnectorStats | None: return stats +class ServeRecordingSecondaryTierManager(MetricsSecondaryTierManager): + """Test-only secondary tier that records service sweeps.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.finished_polls = 0 + self.serve_calls = 0 + + def get_finished_jobs(self) -> Iterable[JobResult]: + self.finished_polls += 1 + return () + + def serve_external_requests(self, parent) -> None: + self.serve_calls += 1 + + +def test_poll_pending_work_services_tiers_between_steps(): + """poll_pending_work() collects finished jobs and serves every tier + without a scheduler step and without consuming the per-step gate.""" + mock_region = _mock_mmap_region(5) + primary = CPUPrimaryTierOffloadingManager(num_chunks=5, mmap_region=mock_region) + tier = ServeRecordingSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="recording", + ) + manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier]) + + manager.poll_pending_work() + manager.poll_pending_work() + assert tier.finished_polls == 2 + assert tier.serve_calls == 2 + assert manager._processed_jobs_this_step is False + + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + manager.on_schedule_end(ctx) + assert tier.finished_polls == 3 + assert tier.serve_calls == 3 + + def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): monkeypatch.setitem( SecondaryTierFactory._registry, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c97374832ffe..bd9c11744fd4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -638,6 +638,18 @@ def has_pending_push_work(self) -> bool: # scheduler alive (e.g. extend has_unfinished_requests). return False + def poll_pending_work(self) -> None: + """Advance scheduler-side work while the engine waits on a model step. + + The engine core calls this at a bounded interval between scheduling + iterations, never concurrently with any other scheduler-side method, + so a connector that services peers from the scheduler side (control + messages, transfer completions) is not limited to one service window + per step. Implementations must not touch scheduler state and must + return promptly when there is nothing to do. + """ + return + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index d291d3852804..858fb69f1ffa 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -588,6 +588,10 @@ def has_pending_block_frees(self) -> bool: def has_pending_push_work(self) -> bool: return any(c.has_pending_push_work() for c in self._connectors) + def poll_pending_work(self) -> None: + for c in self._connectors: + c.poll_pending_work() + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index b253a2399867..00ccdabfa737 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1839,6 +1839,9 @@ def has_pending_push_work(self) -> bool: """ return bool(self._jobs) or self.manager.has_pending_work() + def poll_pending_work(self) -> None: + self.manager.poll_pending_work() + def update_connector_output(self, connector_output: KVConnectorOutput): """ Update KVConnector state from worker-side connectors output. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 73570881753d..6385099cf222 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -202,6 +202,10 @@ def has_pending_push_work(self) -> bool: assert self.connector_scheduler is not None return self.connector_scheduler.has_pending_push_work() + def poll_pending_work(self) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.poll_pending_work() + def update_connector_output(self, connector_output: KVConnectorOutput): assert self.connector_scheduler is not None self.connector_scheduler.update_connector_output(connector_output) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b3d94de3e7db..915f4745960e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -27,6 +27,7 @@ cleanup_dist_env_and_memory, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1 from vllm.envs import enable_envs_cache from vllm.logger import init_logger from vllm.logging_utils.dump_input import dump_engine_exception @@ -102,6 +103,73 @@ _R = TypeVar("_R") # Return type for collective_rpc +# Interval between connector service sweeps while the engine core waits +# for a model step. +KV_CONNECTOR_POLL_INTERVAL_S = 0.005 + + +class ConnectorPoller: + """Services a KV connector's scheduler-side work while the engine core + thread waits on a model step. + + The connector is touched by one thread at a time: the engine core + thread outside ``wait`` and the poller thread inside it. ``wait`` hands + the connector over by raising ``_waiting`` under the lock and takes it + back the same way, so a sweep in progress when the model output arrives + completes before the engine core continues. An exception raised by a + sweep is re-raised from ``wait`` on the engine core thread. + """ + + def __init__( + self, + connector: KVConnectorBase_V1, + interval_s: float = KV_CONNECTOR_POLL_INTERVAL_S, + ) -> None: + self._connector = connector + self._interval_s = interval_s + self._lock = threading.Lock() + self._waiting = False + self._wake = threading.Event() + self._closed = False + self._error: BaseException | None = None + self._thread = threading.Thread( + target=self._run, name="kv-connector-poller", daemon=True + ) + self._thread.start() + + def wait(self, future: Future[Any]) -> Any: + with self._lock: + self._waiting = True + self._wake.set() + try: + return future.result() + finally: + with self._lock: + self._waiting = False + error, self._error = self._error, None + self._wake.clear() + if error is not None: + raise error + + def close(self) -> None: + self._closed = True + self._wake.set() + self._thread.join(timeout=1.0) + + def _run(self) -> None: + while not self._closed: + self._wake.wait() + if self._closed: + return + with self._lock: + if self._waiting and self._error is None: + try: + self._connector.poll_pending_work() + except BaseException as exc: + self._error = exc + time.sleep(self._interval_s) + + class EngineCore: """Inner loop of vLLM's Engine.""" @@ -187,6 +255,9 @@ def __init__( # handshake metadata from all workers so the connector in the scheduler # will have the full context kv_connector = self.scheduler.get_kv_connector() + self.connector_poller: ConnectorPoller | None = ( + ConnectorPoller(kv_connector) if kv_connector is not None else None + ) if kv_connector is not None: # Collect and store KV connector xfer metadata from workers # (after KV cache registration) @@ -620,7 +691,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), ): - model_output = future.result() + model_output = self._wait_for_model_output(future) if model_output is None: model_output = self.model_executor.sample_tokens(grammar_output) @@ -634,6 +705,11 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 + def _wait_for_model_output(self, future: Future[Any]) -> Any: + if self.connector_poller is None: + return future.result() + return self.connector_poller.wait(future) + def post_step(self, model_executed: bool) -> None: # When using async scheduling we can't get draft token ids in advance, # so we update draft token ids in the worker process and don't @@ -719,7 +795,7 @@ def step_with_batch_queue( self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), ): - model_output = future.result() + model_output = self._wait_for_model_output(future) if model_output is None: # None from sample_tokens() implies that the original execute_model() # call failed - raise that exception. @@ -771,6 +847,8 @@ def _process_aborts_queue(self): def shutdown(self): logger.debug_once("[shutdown] EngineCore: tearing down local resources") + if self.connector_poller is not None: + self.connector_poller.close() self.structured_output_manager.clear_backend() if self.model_executor: self.model_executor.shutdown() diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 988747e06e1e..763aa84421ba 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -396,6 +396,15 @@ def has_pending_work(self) -> bool: """ return False + def poll_pending_work(self) -> None: + """Service background work between scheduler steps. + + Called on the scheduler side while the engine waits for a model + step, never concurrently with the per-step methods. The default + does nothing. + """ + return + def reset_cache(self) -> None: """Evict all tracked blocks and reset internal state.""" return diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f19616828daa..6bce2493e0ab 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -852,6 +852,19 @@ def has_pending_work(self) -> bool: or any(tier.has_pending_work() for tier in self.secondary_tiers) ) + @override + def poll_pending_work(self) -> None: + """Service the secondary tiers between scheduler steps. + + The same sweep on_schedule_end() runs at the end of a step: finished + jobs are collected and every tier serves its external requests. The + per-step gate is untouched, so the next step still polls once before + its own lookups. + """ + self._process_finished_jobs() + for tier in self.secondary_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + @override def take_events(self) -> Iterable[OffloadingEvent]: """Yield events owned by the primary and secondary tiers. diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index ffe609d71c1e..d47f55584965 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -75,6 +75,7 @@ import hashlib import json from abc import ABC, abstractmethod +from concurrent.futures import Future from collections.abc import Iterable, Sequence from typing import Literal, NamedTuple @@ -177,6 +178,34 @@ def add_remote_peer( """ ... + 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. + + Returns a future that completes once the peer is usable by + write_blocks(); a failed registration surfaces as the future's + exception. This default runs add_remote_peer() inline and returns + a completed future. Transports whose registration is slow enough + to hold up a scheduling iteration override it to run the work on + a worker thread. + """ + future: Future[None] = Future() + try: + self.add_remote_peer( + peer_id, agent_metadata, base_addr, num_blocks, 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 bf3adf4b604b..ef237e764501 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -6,8 +6,11 @@ from __future__ import annotations +import contextlib import itertools +import sys from collections.abc import Iterable +from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, NamedTuple import numpy as np @@ -23,6 +26,21 @@ logger = init_logger(__name__) + +def _strict_thread_sync() -> Any: + """The agent sync mode that serializes NIXL API calls across threads. + + Registration runs on a worker thread while the engine thread submits + and polls transfers, so the agent must serialize its API. Returns None + when the installed NIXL exposes no sync-mode enum; the agent then keeps + its default and registration still runs on the worker. + """ + module = sys.modules.get(getattr(_NixlAgentConfig, "__module__", "")) + enum = getattr(module, "nixl_thread_sync_t", None) + if enum is None: + return None + return enum.NIXL_THREAD_SYNC_STRICT + # Shared sentinel returned by poll() in the steady state (no inflight, or # no transfer changed state since the last poll). Tuples make it immutable; # callers only iterate / membership-test / equality-check. @@ -68,6 +86,12 @@ def __init__( # transfer_id → _Inflight(peer_id, handle). self._inflight: dict[int, _Inflight] = {} self._next_id = itertools.count() + # peer_id → registration future while add_remote_peer_async runs + # on the worker thread; the agent runs in strict sync mode so the + # worker's registration calls and the caller's transfer calls may + # overlap. + self._registrations: dict[str, Future[None]] = {} + self._registration_executor: ThreadPoolExecutor | None = None self._init(view) @@ -79,9 +103,14 @@ def _init(self, view: memoryview) -> None: if _NixlAgent is None: return + sync_mode = _strict_thread_sync() non_ucx_backends = [b for b in self._backends if b != "UCX"] if non_ucx_backends: - cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True) + cfg = _NixlAgentConfig( + backends=self._backends, + capture_telemetry=True, + sync_mode=sync_mode, + ) logger.info( "NixlTransport %s: NIXL backends=%s", self._agent_name, @@ -89,7 +118,9 @@ def _init(self, view: memoryview) -> None: ) else: cfg = _NixlAgentConfig( - num_threads=self._num_threads, capture_telemetry=True + num_threads=self._num_threads, + capture_telemetry=True, + sync_mode=sync_mode, ) logger.info( "NixlTransport %s: NIXL backends=[UCX] num_threads=%d", @@ -137,7 +168,49 @@ def add_remote_peer( self._peer_nixl_names[peer_id] = nixl_name self._remote_dlists[peer_id] = 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]: + """Run add_remote_peer() on the registration worker. + + Loading the peer's agent metadata and preparing its descriptor list + scale with the peer's block count, so the caller must not hold a + scheduling iteration for them. Peers register one at a time in + arrival order. + """ + if self._registration_executor is None: + self._registration_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="p2p-peer-registration" + ) + future = self._registration_executor.submit( + self.add_remote_peer, + peer_id, + agent_metadata, + base_addr, + num_blocks, + block_len, + ) + self._registrations[peer_id] = future + future.add_done_callback( + lambda f, peer_id=peer_id: self._registrations.pop(peer_id, None) + ) + if future.done(): + # Completed before it was recorded; the callback found nothing. + self._registrations.pop(peer_id, None) + return future + def remove_remote_peer(self, peer_id: str) -> None: + pending = self._registrations.pop(peer_id, None) + if pending is not None: + # A registration still in flight would otherwise complete after + # the removal and leave the peer registered with the agent. + with contextlib.suppress(Exception): + pending.result() nixl_name = self._peer_nixl_names.pop(peer_id, None) dlist = self._remote_dlists.pop(peer_id, None) if self._agent is not None: @@ -292,8 +365,11 @@ def close(self) -> None: return self._release_handles([entry.handle for entry in self._inflight.values()]) self._inflight.clear() - for peer_id in list(self._remote_dlists): + for peer_id in list(self._registrations) + list(self._remote_dlists): self.remove_remote_peer(peer_id) + if self._registration_executor is not None: + self._registration_executor.shutdown(wait=True) + self._registration_executor = None if self._local_dlist is not None: self._agent.release_dlist_handle(self._local_dlist) self._local_dlist = None diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index ba1973b76784..8453518bffca 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -201,8 +201,10 @@ class P2PSecondaryTierManager(SecondaryTierManager): blocks from the peer) and server-role (serving blocks to the peer) over the same control connection. - Single-threaded: every public method runs on the scheduler thread, and - the engine drives polling via ``get_finished_jobs()`` once per step. + Never re-entered: every public method runs on the engine core thread + or, while the engine waits on a model step, on the connector poller, + and the two never overlap. The engine drives polling via + ``get_finished_jobs()`` at least once per step. ``has_pending_work()`` keeps the engine ticking so the control transport and existing sessions are polled even when no requests are scheduled. """ diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index c78918323266..9883856cd7d1 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -449,8 +449,8 @@ def on_lookup( def serve_external_requests(self, parent: ParentManager) -> None: """Resolve inbound peer lookups against the tiering manager. - Called once per scheduler step with a ``parent`` handle valid - only for this call. Drains newly-enqueued LookupMsgs, re-polls + Called at least once per scheduler step with a ``parent`` handle + valid only for this call. Drains newly-enqueued LookupMsgs, re-polls any parked HIT_PENDING / RETRY keys, and releases the bookkeeping for lookups closed since the last serve. """ diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index ba2aca6438c6..5c262eb4cbbb 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -44,6 +44,8 @@ ) if TYPE_CHECKING: + from concurrent.futures import Future + from vllm.v1.kv_offload.base import ReqContext from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport @@ -125,6 +127,10 @@ def __init__( self._send_ready = False # True after the peer acked our ConnectMsg # Msgs waiting to be sent on connection establishment self._queued: list[dict] = [] + # Registration of the peer's NIXL metadata, started by its ConnectMsg. + # ConnectAck goes out once it completes, so the peer never sends a + # fetch this side cannot serve yet. + self._pending_registration: Future[None] | None = None # Consecutive non-protocol dispatch errors. Reset on success. self._dispatch_error_count: int = 0 @@ -166,8 +172,13 @@ def ready(self) -> bool: @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 + """True while inbound loads, outbound transfers, or a peer + registration are outstanding.""" + return ( + self._client.has_active_loads + or self._server.has_inflight_transfers + or self._pending_registration is not None + ) # ------------------------------------------------------------------ # Connection lifecycle @@ -260,6 +271,7 @@ def poll(self) -> SessionPollResult: for msg in self._conn.recv(): self._on_message(msg) + self._maybe_finish_registration() loads = self._client.collect_results() stores = self._server.collect_results() @@ -285,6 +297,7 @@ def close(self) -> SessionCloseResult: ``parent.on_request_finished`` (the manager flushes these on its next ``serve_external_requests``). """ + self._pending_registration = None client_result = self._client.close() failed_stores, failed_serves = self._server.close() @@ -463,19 +476,36 @@ def _on_connect(self, msg: dict) -> None: f"local={self._local_hash_seed!r}. Ensure PYTHONHASHSEED " "(if set) matches on all P2P peers." ) - self._transport.add_remote_peer( - self.peer_id, - agent_metadata=msg[ConnectMsg.AGENT_METADATA], - base_addr=msg[ConnectMsg.BASE_ADDR], - num_blocks=msg[ConnectMsg.NUM_BLOCKS], - block_len=msg[ConnectMsg.BLOCK_LEN], - ) except ValueError as exc: logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc) if self._conn is not None: self._conn.mark_dead() return + self._pending_registration = self._transport.add_remote_peer_async( + self.peer_id, + agent_metadata=msg[ConnectMsg.AGENT_METADATA], + base_addr=msg[ConnectMsg.BASE_ADDR], + num_blocks=msg[ConnectMsg.NUM_BLOCKS], + block_len=msg[ConnectMsg.BLOCK_LEN], + ) + + def _maybe_finish_registration(self) -> None: + """Send ConnectAck once the peer's registration has completed. + + A failed registration rejects the peer the same way a ConnectMsg + validation failure does. + """ + future = self._pending_registration + if future is None or not future.done(): + return + self._pending_registration = None + exc = future.exception() + if exc is not None: + logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc) + if self._conn is not None: + self._conn.mark_dead() + return if self._conn is not None: self._conn.send( {