[NIXL] Drop unreachable producer-side notif counting in push mode (#62210)
Signed-off-by: David Cheung <d7cheung@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: vllm-agent <claw@inferact.ai>
This commit is contained in:
David Cheung
co-authored by
Claude Opus 5.5
vllm-agent
Claude Opus 5.5
vllm-agent
parent
255f14599e
commit
7879f3687b
@@ -929,14 +929,8 @@ class TestPushWriterNotifs:
|
||||
notif_msg = f"{request_id}:1".encode()
|
||||
w._pending_completion_notifs.put(notif_msg)
|
||||
|
||||
# transfer_topo is consulted only for the producer-side path; we
|
||||
# make it a MagicMock because the D-side branch returns early.
|
||||
w.transfer_topo = MagicMock()
|
||||
w._get_new_notifs()
|
||||
|
||||
notified = w._get_new_notifs()
|
||||
|
||||
# Notif consumed; D-side just touches _recving_transfers.
|
||||
assert notified == set()
|
||||
assert request_id in w._recving_transfers
|
||||
|
||||
def test_get_transfer_results_evicts_completed_state(self):
|
||||
@@ -1352,12 +1346,10 @@ class TestPushWriterNegative:
|
||||
logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"),
|
||||
)
|
||||
w = _StubWriterWorker.fresh()
|
||||
w.transfer_topo = MagicMock()
|
||||
# Forward a completion notif for an unknown request_id.
|
||||
w._pending_completion_notifs.put(b"never-heard-of-you:1")
|
||||
|
||||
notified = w._get_new_notifs()
|
||||
assert notified == set()
|
||||
w._get_new_notifs()
|
||||
# Did not register anywhere.
|
||||
assert "never-heard-of-you" not in w._recving_transfers
|
||||
assert w.xfer_stats.data["num_notifications_after_expiry"] == [1]
|
||||
@@ -1381,7 +1373,6 @@ class TestPushWriterNegative:
|
||||
leases of tracked P-side requests on the engine main thread, and
|
||||
ignore request IDs that aren't being tracked."""
|
||||
w = _StubWriterWorker.fresh()
|
||||
w.transfer_topo = MagicMock()
|
||||
# _handle_heartbeat reads ``self._lease_extension`` (set in the
|
||||
# real ``__init__``).
|
||||
w._lease_extension = 10
|
||||
@@ -1395,8 +1386,7 @@ class TestPushWriterNegative:
|
||||
# and another tracked one.
|
||||
w._pending_completion_notifs.put(b"HB:req-a,req-unknown,req-b")
|
||||
|
||||
notified = w._get_new_notifs()
|
||||
assert notified == set()
|
||||
w._get_new_notifs()
|
||||
|
||||
# Tracked leases were renewed strictly forward in time.
|
||||
now = time.perf_counter()
|
||||
@@ -1421,20 +1411,19 @@ class TestPushPipelineParallel:
|
||||
"""Each PP stage WRITEs its own layers and sends one notif; D must
|
||||
collect pp_size notifs before reporting the recv done."""
|
||||
w = _StubWriterWorker.fresh()
|
||||
w.transfer_topo = MagicMock()
|
||||
request_id = "req-pp-2"
|
||||
w._recving_metadata[request_id] = MagicMock(pp_size=2)
|
||||
notif = f"{request_id}:1".encode()
|
||||
|
||||
# First stage: counted, not yet done.
|
||||
w._pending_completion_notifs.put(notif)
|
||||
assert w._get_new_notifs() == set()
|
||||
w._get_new_notifs()
|
||||
assert request_id not in w._recving_transfers
|
||||
assert w.consumer_notification_counts_by_req[request_id] == 1
|
||||
|
||||
# Second (final) stage: now reported done.
|
||||
w._pending_completion_notifs.put(notif)
|
||||
assert w._get_new_notifs() == set()
|
||||
w._get_new_notifs()
|
||||
assert request_id in w._recving_transfers
|
||||
assert request_id not in w.consumer_notification_counts_by_req
|
||||
|
||||
|
||||
@@ -748,11 +748,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
|
||||
"""Drain HB / completion notifs forwarded by the writer thread.
|
||||
|
||||
The writer owns ``nixl_wrapper.get_new_notifs`` for push; PUSH_REG
|
||||
notifs are handled there. Everything else is forwarded here for
|
||||
existing accounting.
|
||||
notifs are handled there. P learns of finished pushes from its WRITE
|
||||
handles, so this never reports finished sends.
|
||||
"""
|
||||
assert self.transfer_topo is not None
|
||||
notified_req_ids: set[str] = set()
|
||||
while True:
|
||||
try:
|
||||
notif = self._pending_completion_notifs.get_nowait()
|
||||
@@ -766,46 +764,27 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
|
||||
|
||||
req_id, tp_size = msg.rsplit(":", 1)
|
||||
|
||||
# Not tracked as a P-side send/process for this notif.
|
||||
if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process:
|
||||
if (meta := self._recving_metadata.get(req_id)) is not None:
|
||||
if (meta := self._recving_metadata.get(req_id)) is None:
|
||||
# The lease may have expired before the notif arrived.
|
||||
self.xfer_stats.record_notification_after_expiry()
|
||||
logger.error(
|
||||
"Unrecognized request %s notif (may have expired).", req_id
|
||||
)
|
||||
continue
|
||||
# Consumer waits for one notif per producer rank writing
|
||||
# here: pp_size stages * producers-per-consumer (>1 when
|
||||
# producer TP > consumer TP; tp_size is the producer TP).
|
||||
producers_per_consumer = max(1, int(tp_size) // self.world_size)
|
||||
expected_notifs = meta.pp_size * producers_per_consumer
|
||||
self.consumer_notification_counts_by_req[req_id] += 1
|
||||
notifs = self.consumer_notification_counts_by_req[req_id]
|
||||
if notifs < expected_notifs:
|
||||
if self.consumer_notification_counts_by_req[req_id] < expected_notifs:
|
||||
continue
|
||||
del self.consumer_notification_counts_by_req[req_id]
|
||||
# P drove the transfer (we own no NIXL handle), so
|
||||
# materialise an empty ``_recving_transfers`` entry for
|
||||
# ``_pop_done_transfers`` to report done.
|
||||
self._recving_transfers.setdefault(req_id, [])
|
||||
else:
|
||||
# Not tracked on either side (lease may have expired
|
||||
# before the notif arrived). Log and skip.
|
||||
self.xfer_stats.record_notification_after_expiry()
|
||||
logger.error(
|
||||
"Unrecognized request %s notif (may have expired).",
|
||||
req_id,
|
||||
)
|
||||
continue
|
||||
|
||||
n_consumers = int(tp_size)
|
||||
tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
|
||||
consumers_per_producer = -tp_ratio if n_consumers > self.world_size else 1
|
||||
self.consumer_notification_counts_by_req[req_id] += 1
|
||||
if (
|
||||
self.consumer_notification_counts_by_req[req_id]
|
||||
== consumers_per_producer
|
||||
):
|
||||
notified_req_ids.add(req_id)
|
||||
del self.consumer_notification_counts_by_req[req_id]
|
||||
self._reqs_to_process.remove(req_id)
|
||||
self._reqs_to_send.pop(req_id, None)
|
||||
return notified_req_ids
|
||||
return set()
|
||||
|
||||
def get_transfer_results(self) -> KVConnectorTransferResults:
|
||||
# Engine main thread asking for completions: also wake the writer
|
||||
@@ -835,7 +814,6 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
|
||||
for req_id in done_pushing:
|
||||
self._reqs_to_send.pop(req_id, None)
|
||||
self._reqs_to_process.discard(req_id)
|
||||
self.consumer_notification_counts_by_req.pop(req_id, None)
|
||||
done_sending.add(req_id)
|
||||
|
||||
# Tell the writer to drop any state it still holds for any
|
||||
|
||||
Reference in New Issue
Block a user