iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV Offload] Bound unanswered P2P lookup waits (#56836)

Browse Source 
Signed-off-by: jacklin78911-collab <jacklin78911@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.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>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
This commit is contained in:
Liqian Lin
2026-10-15 08:23:34 +03:00
committed by GitHub
co-authored by Codex Robert Shaw Itay Etelis mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Or Ozeri Kevin H. Luu
parent 81e1c3f183
commit 3f1bcfd82b
6 changed files with 386 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_offload/tiering/p2p/conftest.py
+15
View File
@@ -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
tests/v1/kv_offload/tiering/p2p/test_manager.py
+103
View File
@@ -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()
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+216 -7
View File
@@ -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 == []
# ---------------------------------------------------------------------------
vllm/v1/kv_offload/tiering/p2p/session/client.py
+30 -10
View File
@@ -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)
vllm/v1/kv_offload/tiering/p2p/session/server.py
+5 -4
View File
@@ -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]:
vllm/v1/kv_offload/tiering/p2p/session/session.py
+17 -3
View File
@@ -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(