diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index d92857c341..e2f9e99b49 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 6d82fbcc0a..38e7d0fbb7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -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: - # 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: - 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, - ) + 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 - - 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 + # 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 - 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 + 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, []) + 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