diff --git a/tests/v1/kv_offload/tiering/p2p/conftest.py b/tests/v1/kv_offload/tiering/p2p/conftest.py new file mode 100644 index 0000000000..e5b27f8cf7 --- /dev/null +++ b/tests/v1/kv_offload/tiering/p2p/conftest.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest + +from vllm.v1.kv_offload.tiering.p2p.session import client as client_module + + +@pytest.fixture +def lookup_clock(monkeypatch): + """Advance monotonic time for P2P deadlines without sleeping.""" + clock = SimpleNamespace(now=0.0) + monkeypatch.setattr(client_module.time, "monotonic", lambda: clock.now) + return clock diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 0ca2735005..2282171fdb 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -37,6 +37,7 @@ from vllm.v1.kv_offload.tiering.p2p.session import ( SessionPollResult, StoreResult, ) +from vllm.v1.kv_offload.tiering.p2p.session.client import _LOOKUP_TIMEOUT_S # --------------------------------------------------------------------------- # Helpers @@ -1172,6 +1173,108 @@ def _build_paired_managers() -> tuple[P2PSecondaryTierManager, P2PSecondaryTierM class TestBidirectionalManager: """Two managers each load FROM and serve TO the other over a single peer.""" + def test_unanswered_lookup_falls_back_to_local_prefill(self, lookup_clock): + mgr_a, mgr_b = _build_paired_managers() + ctx = _req_context(_remote_kv_source_kv_params("B", 2)) + mgr_a.on_new_request(ctx) + for _ in range(3): + list(mgr_a.get_finished_jobs()) + list(mgr_b.get_finished_jobs()) + assert mgr_a._sessions["B:2"].ready + assert mgr_a.lookup(b"key", ctx) == LookupResult.RETRY + mgr_a.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=[])) + + # Stop the server's progress without closing its live connection. + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + assert list(mgr_a.get_finished_jobs()) == [] + assert mgr_a.lookup(b"key", ctx) == LookupResult.RETRY + lookup_clock.now = _LOOKUP_TIMEOUT_S + assert list(mgr_a.get_finished_jobs()) == [] + assert mgr_a.lookup(b"key", ctx) == LookupResult.MISS + assert "B:2" not in mgr_a._sessions + assert "B:2" not in mgr_a._data._remote_peers + + # A replacement session must not restart the expired request's probes. + mgr_a._sessions["B:2"] = _FakeSession(peer_id="B:2") # type: ignore[assignment] + assert mgr_a.lookup(b"key", ctx) == LookupResult.MISS + mgr_a.on_request_finished(ctx) + assert "req-1" not in mgr_a._failed_req_ids + + def test_lookup_expiry_allows_zmq_reconnect(self, lookup_clock): + """The peer must retire the timed-out session before a new handshake.""" + import zmq + from zmq.utils.monitor import recv_monitor_message + + from .test_zmq_transport import _make_transport + + managers = [] + monitors = [] + try: + for _ in range(2): + mgr = _make_manager() + mgr._control, port = _make_transport() + mgr._local_id = f"127.0.0.1:{port}" + mgr._data = _FakeData(mgr._local_id) + managers.append(mgr) + monitors.append( + mgr._control._router.get_monitor_socket(zmq.EVENT_DISCONNECTED) + ) + mgr_a, mgr_b = managers + peer_a, peer_b = mgr_a._local_id, mgr_b._local_id + port_b = int(peer_b.rsplit(":", 1)[1]) + + def poll_until(condition): + deadline = time.perf_counter() + 5.0 + while time.perf_counter() < deadline: + list(mgr_a.get_finished_jobs()) + list(mgr_b.get_finished_jobs()) + if condition(): + return + time.sleep(0.005) + pytest.fail("P2P session did not reach the expected state") + + def both_ready(): + return ( + peer_b in mgr_a._sessions + and peer_a in mgr_b._sessions + and mgr_a._sessions[peer_b].ready + and mgr_b._sessions[peer_a].ready + ) + + ctx = _req_context(_remote_kv_source_kv_params("127.0.0.1", port_b)) + mgr_a.on_new_request(ctx) + poll_until(both_ready) + old_peer_session = mgr_b._sessions[peer_a] + assert mgr_a.lookup(b"key", ctx) == LookupResult.RETRY + mgr_a.on_schedule_end( + ScheduleEndContext(new_req_ids=[], preempted_req_ids=[]) + ) + + lookup_clock.now = _LOOKUP_TIMEOUT_S + list(mgr_a.get_finished_jobs()) + assert mgr_a.lookup(b"key", ctx) == LookupResult.MISS + assert peer_b not in mgr_a._sessions + poll_until(lambda: peer_a not in mgr_b._sessions) + # Wait for old TCP identities to retire before reconnecting. + for monitor in monitors: + assert monitor.poll(5000), "old TCP connection did not disconnect" + assert recv_monitor_message(monitor)["event"] == zmq.EVENT_DISCONNECTED + + new_ctx = _req_context( + _remote_kv_source_kv_params("127.0.0.1", port_b, "req-2") + ) + mgr_a.on_new_request(new_ctx) + poll_until(both_ready) + assert mgr_b._sessions[peer_a] is not old_peer_session + assert mgr_a.lookup(b"key", ctx) == LookupResult.MISS + finally: + for monitor in monitors: + monitor.close(linger=0) + # Avoid queuing shutdown messages to an already closed peer. + for mgr in managers: + mgr._control.close() + mgr.shutdown() + def test_both_loads_succeed(self): mgr_a, mgr_b = _build_paired_managers() diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 63abc642b0..89f3ac0cd6 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -36,6 +36,7 @@ from vllm.v1.kv_offload.tiering.p2p.session import ( from vllm.v1.kv_offload.tiering.p2p.session.client import ( _ABORT_ACK_TIMEOUT_S, _LOAD_TIMEOUT_S, + _LOOKUP_TIMEOUT_S, ) from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, @@ -839,6 +840,215 @@ class TestClientFlows: class TestLookupFlow: """Consumer-side state machine for do_p2p_fetch lookups.""" + def test_lookup_deadline_starts_after_handshake(self, lookup_clock): + session, conn, _ = _make_session() + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + # Complete the handshake before its independent timeout expires. + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + session.poll() + assert session.alive + + _activate(session, conn) + session.flush_pending_lookups() + lookups = [m for m in conn._sent if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(lookups) == 1 + lookup_clock.now += _LOOKUP_TIMEOUT_S - 1 + session.poll() + assert session.alive + lookup_clock.now += 1 + session.poll() + assert not session.alive + assert session.close().failed_req_ids == ["req-1"] + + def test_partial_response_rearms_lookup_deadline(self, lookup_clock): + session, conn, _ = _make_session() + _activate(session, conn) + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is None + session.register_lookup("req-1", b"hC") + session.flush_pending_lookups() + + lookup_clock.now += 1 + session.poll() + assert session.alive + lookup_clock.now += _LOOKUP_TIMEOUT_S - 1 + session.poll() + assert not session.alive + assert session.close().failed_req_ids == ["req-1"] + + def test_lookup_progress_keeps_overlapping_batches_alive(self, lookup_clock): + """Prompt batch replies keep a long-running lookup phase alive.""" + session, conn, _ = _make_session() + _activate(session, conn) + previous_key = b"key-0" + session.register_lookup("req-1", previous_key) + session.flush_pending_lookups() + + for step in range(1, int(2 * _LOOKUP_TIMEOUT_S) + 1): + lookup_clock.now = float(step) + next_key = f"key-{step}".encode() + hit = step % 2 == 0 + session.register_lookup("req-1", next_key) + session.flush_pending_lookups() + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [previous_key], + LookupRespMsg.HITS: [hit], + } + ) + session.poll() + assert session.alive + assert session.register_lookup("req-1", previous_key) is hit + previous_key = next_key + + def test_buffered_lookup_response_precedes_expiry(self, lookup_clock): + session, conn, _ = _make_session() + _activate(session, conn) + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + lookup_clock.now = _LOOKUP_TIMEOUT_S + session.poll() + assert session.alive + assert session.register_lookup("req-1", b"hA") is True + + def test_duplicate_responses_and_new_keys_do_not_rearm(self, lookup_clock): + session, conn, _ = _make_session() + _activate(session, conn) + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + response = { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA", b"unknown"], + LookupRespMsg.HITS: [True, False], + } + conn.enqueue(response) + session.poll() + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + conn.enqueue(response) + session.register_lookup("req-1", b"hC") + session.flush_pending_lookups() + session.poll() + assert session.alive + lookup_clock.now = _LOOKUP_TIMEOUT_S + session.poll() + assert not session.alive + + @pytest.mark.parametrize("partial_fetch", [False, True]) + def test_lookup_expiry_preserves_completed_transfers( + self, lookup_clock, partial_fetch + ): + session, conn, transport = _make_session() + _activate(session, conn) + session.register_lookup("stalled", b"key") + session.flush_pending_lookups() + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + session.request_blocks(1, "load", [b"key"], [0]) + session.add_stored_blocks("store", [b"key"], [1], 2) + if partial_fetch: + session.add_stored_blocks("store", [b"unused"], [2], 3) + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: "store", + FetchMsg.KEYS: [b"key"], + FetchMsg.BLOCK_INDEXES: [0], + FetchMsg.ROUND_SEQ: 0, + } + ) + session.poll() + transport._poll_done = list(transport._transfers) + conn.enqueue( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.KV_REQUEST_ID: "load", + TransferDoneMsg.SUCCESS: True, + TransferDoneMsg.ROUND_SEQ: 0, + } + ) + lookup_clock.now = _LOOKUP_TIMEOUT_S + result = session.poll() + assert not session.alive + assert result.loads == [LoadResult(1, "load", True)] + expected_stores = [StoreResult(2, True)] + if partial_fetch: + expected_stores.append(StoreResult(3, True)) + assert result.stores == expected_stores + closed = session.close() + assert closed.failed_req_ids == ["stalled"] + assert closed.failed_jobs == [] + assert closed.failed_stores == [] + + @pytest.mark.parametrize("action", ["resolve", "finish", "fetch"]) + def test_lookup_deadline_cleared_before_next_batch(self, lookup_clock, action): + session, conn, _ = _make_session() + _activate(session, conn) + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + lookup_clock.now = _LOOKUP_TIMEOUT_S - 1 + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + if action == "resolve": + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hB"], + LookupRespMsg.HITS: [False], + } + ) + session.poll() + elif action == "finish": + session.finish_request("req-1") + else: + session.request_blocks(1, "req-1", [b"hA"], [0]) + + lookup_clock.now += 2 + session.poll() + assert session.alive + session.register_lookup("req-1", b"hC") + session.flush_pending_lookups() + lookup_clock.now += _LOOKUP_TIMEOUT_S - 1 + session.poll() + assert session.alive + lookup_clock.now += 1 + session.poll() + assert not session.alive + def test_aggregate_flush_resolve_round_trip(self): """register_lookup → flush sends one LookupMsg → response resolves entries → register_lookup returns the cached bool on @@ -2531,20 +2741,19 @@ class TestFinishRequestServerSide: # Transfer 1 completes — now ``_has_inflight_for("req-1")`` is False # and the elif branch in collect_results fires _finalize(success=False). - # The k1 success result is direct; the k2 failure result is queued - # in _pending_store_results and surfaces on the NEXT poll. + # Both the k1 success and the deferred k2 failure surface in this poll. transport._poll_done.append(tid_1) - stores_first = session.poll().stores - assert StoreResult(job_id=100, success=True) in stores_first + assert session.poll().stores == [ + StoreResult(job_id=100, success=True), + StoreResult(job_id=200, success=False), + ] # Outbound state cleaned up; peer notified with success=False. assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is False - # Next poll drains the queued failure. - stores_second = session.poll().stores - assert StoreResult(job_id=200, success=False) in stores_second + assert session.poll().stores == [] # --------------------------------------------------------------------------- diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index 78eb1b76ea..4af2047a7b 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -32,6 +32,7 @@ logger = init_logger(__name__) _LOAD_TIMEOUT_S = 30.0 _ABORT_ACK_TIMEOUT_S = 10.0 +_LOOKUP_TIMEOUT_S = 30.0 @dataclass @@ -59,10 +60,8 @@ class _ClientRequestState: # -- Lookup phase (symmetric P2P only; untouched for PD) -- # Probe outcome per OffloadKey: None while in-flight (registered/sent - # but unresolved), True/False once a LookupRespMsg lands. There is no - # timeout — finish (via finish_request) is guaranteed after - # the request's lookup() calls and clears every probe, so an - # unanswered probe simply stays None until then. + # but unresolved), True/False once a LookupRespMsg lands. Unanswered + # probes expire the session so the manager can fall back to local compute. probes: dict[OffloadKey, bool | None] = field(default_factory=dict) # OffloadKeys registered but not yet flushed onto the wire. Drained and # cleared by the next flush_pending_lookups. @@ -136,6 +135,9 @@ class ClientRole: # Kept in exact sync with ``st.loads``. self._active_loads: set[str] = set() self._completed_loads: list[LoadResult] = [] + # Per-request progress deadlines. Only newly resolved probes extend + # them; additional keys and repeated scheduler polls do not. + self._lookup_deadlines: dict[str, float] = {} # ------------------------------------------------------------------ # State helpers @@ -172,6 +174,14 @@ class ClientRole: """True if any kv_request_id has a fetch in flight.""" return bool(self._active_loads) + def get_expired_lookup(self) -> tuple[str, float] | None: + """Return an expired request and seconds without progress, if any.""" + now = time.monotonic() + for req_id, deadline in self._lookup_deadlines.items(): + if now >= deadline: + return req_id, now - deadline + _LOOKUP_TIMEOUT_S + return None + # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ @@ -228,6 +238,7 @@ class ClientRole: ) assert all(st.probes.get(key) is True for key in keys) st.probes.clear() + self._lookup_deadlines.pop(kv_request_id, None) def finish(self, kv_request_id: str) -> None: """Finish a request: abort in-flight loads and release lookup state. @@ -272,6 +283,7 @@ class ClientRole: ) st.probes.clear() st.unsent.clear() + self._lookup_deadlines.pop(kv_request_id, None) self._flush_pending.discard(kv_request_id) self._maybe_prune(kv_request_id) @@ -377,7 +389,7 @@ class ClientRole: ) return None - def flush_pending_lookups(self) -> None: + def flush_pending_lookups(self, send_ready: bool) -> None: """Send a LookupMsg for each kv_request_id with unsent entries. Called once per scheduler step from the manager's @@ -387,13 +399,14 @@ class ClientRole: keys, all tagged with the current round. register_lookup() de-dups in-flight and already-resolved (req_id, key) pairs, so each LookupMsg carries only the keys first probed in that step. - Send-gating is handled by the injected ``_send`` callback - (queues until ConnectAckMsg if needed). + Wait for ConnectAckMsg before sending or starting the response deadline. Only requests that registered new keys since the last flush are visited — the ``_flush_pending`` work-list avoids scanning every live request each scheduler step. """ + if not send_ready: + return for req_id in self._flush_pending: st = self._requests.get(req_id) if st is None or not st.unsent: @@ -401,6 +414,9 @@ class ClientRole: # The peer now holds lookup state for this id; finish owes a # terminal empty FetchMsg until a fetch closes it. st.peer_lookup_open = True + self._lookup_deadlines.setdefault( + req_id, time.monotonic() + _LOOKUP_TIMEOUT_S + ) logger.debug( "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d", self._peer_id, @@ -443,10 +459,16 @@ class ClientRole: st = self._requests.get(kv_request_id) if st is None: return + progressed = False for h, hit in zip(keys, hits): key = OffloadKey(h) if key in st.probes: + progressed |= st.probes[key] is None st.probes[key] = hit + if all(hit is not None for hit in st.probes.values()): + self._lookup_deadlines.pop(kv_request_id, None) + elif progressed: + self._lookup_deadlines[kv_request_id] = time.monotonic() + _LOOKUP_TIMEOUT_S def collect_results(self) -> list[LoadResult]: """Walk load timeouts and drain completed loads. @@ -454,9 +476,6 @@ class ClientRole: Loads past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and enter the aborting phase. Aborting loads past ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed. - - Lookups have no timeout: an unanswered probe stays None (RETRY) - until finish_request clears it — see ``_ClientRequestState.probes``. """ now = time.monotonic() to_remove: list[tuple[str, int]] = [] @@ -525,4 +544,5 @@ class ClientRole: self._flush_pending.clear() self._active_loads.clear() self._completed_loads.clear() + self._lookup_deadlines.clear() return ClientCloseResult(failed_jobs=failed_jobs, failed_req_ids=failed_req_ids) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index fd43e31ec5..6ac7af5ef0 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -755,10 +755,6 @@ class ServerRole: """ results: list[StoreResult] = self._timeout_pending_store_jobs() - if self._pending_store_results: - results.extend(self._pending_store_results) - self._pending_store_results.clear() - # Scope the poll to this peer: the transport is shared across all peer # sessions of the engine, and poll() drains completed handles. An # unscoped poll here would consume sibling sessions' completions and @@ -833,6 +829,11 @@ class ServerRole: self._finalize_cancelled(xfer.kv_request_id, xfer.round_key) self._maybe_prune(xfer.kv_request_id) + # Finalizing transfers can queue results for supplied but unfetched blocks. + if self._pending_store_results: + results.extend(self._pending_store_results) + self._pending_store_results.clear() + return results def collect_idle_timeouts(self) -> list[StoreResult]: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index 6d3d990034..bfb26719d4 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -264,10 +264,9 @@ class P2PSession: """Flush any aggregated symmetric-P2P lookups for this peer. Called once per scheduler step from the manager's - ``on_schedule_end()``. Send-gating is handled inside the - client's ``_send`` callback (queues until ConnectAckMsg). + ``on_schedule_end()``. Probes wait for ConnectAckMsg before sending. """ - self._client.flush_pending_lookups() + self._client.flush_pending_lookups(send_ready=self._send_ready) def serve_external_requests(self, parent: ParentManager) -> None: """Resolve inbound peer lookups against the tiering manager. @@ -329,6 +328,21 @@ class P2PSession: stores = self._server.collect_results() self._server.drain_pending_aborts() + # Drain buffered replies and completed transfers before expiring a peer. + expired_lookup = self._client.get_expired_lookup() + if expired_lookup is not None: + req_id, elapsed = expired_lookup + logger.warning( + "P2PSession %s: lookup response timed out for kv_request_id=%s " + "after %.1fs without progress; disconnecting peer session", + self.peer_id, + req_id, + elapsed, + ) + with contextlib.suppress(Exception): + self._conn.send({TYPE_KEY: DisconnectMsg.TYPE}) + self._conn.mark_dead() + new_fetch_ids = self._new_fetch_ids self._new_fetch_ids = [] return SessionPollResult(