From 802feb6770a0522ce8d2999262d95f4618e04adf Mon Sep 17 00:00:00 2001 From: Akshay Viswanathan Date: Tue, 13 Oct 2026 00:29:31 +0200 Subject: [PATCH] [Bugfix][KVConnector] Fix MoRI-IO WRITE-mode requests stranded in WAITING_FOR_REMOTE_KVS (#57705) Signed-off-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- .../unit/test_moriio_write_done_routing.py | 104 +++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 9 +- .../v1/moriio/moriio_connector.py | 73 ++++++++--- .../kv_connector/v1/moriio/moriio_engine.py | 121 +++++++++++++----- 4 files changed, 255 insertions(+), 52 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index adac007d19..57b6893869 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -3,20 +3,30 @@ """Unit tests for mori-io write_done routing (#51681).""" import threading +import time +from collections import OrderedDict from types import SimpleNamespace from typing import Any +from unittest.mock import patch import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ROLE, + MoRIIOMode, RemoteAllocInfo, ReqMeta, WriteTask, get_port_offset, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorScheduler, MoRIIOConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) from .utils import make_moriio_writer @@ -79,6 +89,7 @@ def _write_task( *, multi_pod_hosts: list[str] | None = None, remote_dp_size_local: int = 0, + remote_dp_size: int = 1, request_id: str = "req", ) -> WriteTask: return WriteTask( @@ -93,9 +104,21 @@ def _write_task( remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts or [], remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) +def _real_wrapper() -> MoRIIOWrapper: + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) + wrapper.lock = threading.Lock() + wrapper.done_req_ids = [] + wrapper.done_remote_allocate_req_dict = {} + wrapper.done_write_cache_req_ids = [] + wrapper.failed_write_cache_req_ids = [] + wrapper._terminal_transfer_ids = OrderedDict() + return wrapper + + def test_execute_write_task_ip_is_per_task_under_overwrite(): alloc = {"tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0)} worker = _fake_worker(alloc) @@ -269,3 +292,84 @@ def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty(): assert captured["remote_dp_size_local"] == 4 assert captured["local_block_ids"] == [1] assert captured["remote_block_ids"] == [2] + + +def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): + wrapper = _real_wrapper() + sent: list[tuple[str, int, str]] = [] + wrapper.send_notify = lambda tid, ip, port, message_type: sent.append( + (ip, port, message_type) + ) + writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=0)) + task = _write_task( + "tA", + "10.0.0.1", + multi_pod_hosts=["10.0.0.1", "10.0.0.2"], + remote_dp_size_local=2, + remote_dp_size=4, + ) + task.enqueue_time = time.perf_counter() - 2 * writer._defer_timeout + writer._deferred_tasks.append(task) + + writer._process_deferred_tasks() + + assert writer._deferred_tasks == [] + assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["tA"] + # Global DP ranks 0-3, two per pod: the owning rank is unknown. + assert sent == [ + (host, 20000 + get_port_offset(local_rank, 0), "write_failed") + for host in ("10.0.0.1", "10.0.0.2") + for local_rank in (0, 1) + ] + + # A late decode allocation must not resurrect a write into freed blocks. + with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER): + wrapper._handle_remote_blocks_message( + {"transfer_id": "tA", "block_notify_list": [7], "decode_rank": 1} + ) + assert wrapper.done_remote_allocate_req_dict == {} + + +def test_write_scheduler_waits_for_worker_ack_past_defer_deadline(): + scheduler = SimpleNamespace( + is_producer=True, + mode=MoRIIOMode.WRITE, + _defer_timeout=60.0, + _deferred_send_deadlines={"req-1": (time.monotonic() - 1.0, "tA")}, + _pending_sent_acks={}, + unmap_request_id=lambda *args, **kwargs: None, + ) + output = SimpleNamespace(finished_sending=set()) + + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending is None + assert "req-1" in scheduler._deferred_send_deadlines + + output.finished_sending = {"req-1"} + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending == {"req-1"} + assert scheduler._deferred_send_deadlines == {} + + +def test_write_failed_is_reported_as_failed_recving_on_owning_rank(): + wrapper = _real_wrapper() + with patch.object(moriio_engine, "get_role", return_value=ROLE.CONSUMER): + for tid in ("tA", "tOtherRank"): + wrapper._handle_structured_message( + {"type": "write_failed", "transfer_id": tid} + ) + worker = SimpleNamespace( + mode=MoRIIOMode.WRITE, + is_producer=False, + moriio_wrapper=wrapper, + transfer_id_to_request_id={"tA": "req-1"}, + _unmatched_write_failures=OrderedDict(), + get_finished=lambda: (set(), set()), + ) + + results = MoRIIOConnectorWorker.get_transfer_results(worker) + + assert results.finished_recving == {"req-1"} + assert results.failed_recving == {"req-1"} + assert wrapper.pop_finished_write_req_ids() == set() + assert list(worker._unmatched_write_failures) == ["tOtherRank"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index ae1eecf6e5..e82d066818 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -71,6 +71,7 @@ class WriteTask: remote_ip: str multi_pod_hosts: list[str] = field(default_factory=list) remote_dp_size_local: int = 0 + remote_dp_size: int = 1 enqueue_time: float = field(default_factory=time.perf_counter) retried: int = 0 @@ -307,8 +308,8 @@ class MoRIIOConfig: # WRITE mode. # transfer_timeout -> Timeout for waiting_for_transfer_complete before # raising TransferError (sec). - # defer_timeout -> Timeout before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed (sec). + # defer_timeout -> Timeout for deferred WRITE allocation or a missing + # READ release notification (sec). # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA # completion never arrived is aborted (sec). @@ -414,8 +415,8 @@ class MoRIIOConstants: # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError. # Overridable via kv_connector_extra_config["transfer_timeout"]. DEFAULT_TRANSFER_TIMEOUT = 30.0 - # Timeout (seconds) before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed. + # Timeout (seconds) for deferred WRITE allocation or a missing READ release + # notification. # Overridable via kv_connector_extra_config["defer_timeout"]. DEFAULT_DEFER_TIMEOUT = 60.0 # Timeout (seconds) before an in-flight recv whose RDMA completion was lost diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 082643778e..c398cb20a8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -5,7 +5,7 @@ import math import queue import threading import time -from collections import defaultdict +from collections import OrderedDict, defaultdict from collections.abc import Collection from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast @@ -25,6 +25,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( @@ -116,6 +117,9 @@ logger = init_logger(__name__) _SQ_FULL_BACKOFF_INITIAL_S = 0.001 _SQ_FULL_BACKOFF_MAX_S = 0.05 _MAX_LOCAL_MAMBA_TAIL_BLOCKS = 1 +# A WRITE failure is broadcast to every decode DP rank, so ranks that never +# see the request would otherwise keep its transfer id forever. +_MAX_UNMATCHED_WRITE_FAILURES = 4096 try: @@ -356,6 +360,12 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + return self.connector_worker.get_transfer_results() + def get_block_ids_with_load_errors(self) -> set[int]: """Blocks whose MoRIIO read failed, for the scheduler to recompute.""" if self.connector_worker is None: @@ -667,12 +677,10 @@ class MoRIIOConnectorScheduler: set_role(ROLE.CONSUMER) # Reqs to send and their expiration time self._reqs_need_send: dict[ReqId, float] = {} - # Deadlines for requests whose block freeing was deferred. - # Survives across scheduler steps. If the worker never reports - # finished_sending before the deadline, we inject them into - # connector_output.finished_sending so the scheduler frees the blocks to avoid - # hanging indefinitely waiting for a free notification that never comes. - # Value: (deadline, transfer_id) for unmapping after mutation. + # Requests whose block freeing is deferred. READ mode uses the deadline + # as a missing-ACK watchdog; WRITE mode waits for the worker ACK because + # only the worker can mark the transfer terminal before blocks are freed. + # Value: (deadline, transfer_id) for unmapping after completion. self._deferred_send_deadlines: dict[ReqId, tuple[float, TransferId | None]] = {} self._defer_timeout = float( self.kv_transfer_config.kv_connector_extra_config.get( @@ -1131,7 +1139,14 @@ class MoRIIOConnectorScheduler: _should_notify = self._global_dp_rank == remote_dp_rank else: _should_notify = self._is_kv_master - if _should_notify: + if _should_notify and num_external_tokens == 0: + # Nothing left to fetch (full local prefix hit), so release + # the producer's blocks instead of sending an empty + # allocation, which it rejects. + self._release_write_prefill_blocks( + request.request_id, request.kv_transfer_params + ) + elif _should_notify: peer_zmq = get_peer_zmq_from_request_id( request.request_id, is_producer=False ) @@ -1159,12 +1174,6 @@ class MoRIIOConnectorScheduler: f"remote_notify_port={remote_notify_port!r})" ) - # num_external_tokens == 0: nothing to push, so don't tell - # the producer to write into these blocks. - block_notify_list = ( - blocks.get_block_ids()[0] if num_external_tokens > 0 else [] - ) - # Wide-EP multi-pod: a pod binds notify sockets only for # its LOCAL ranks, so the port offset must use the per-pod # local rank (% dp_local), not the global rank. Single-pod @@ -1191,7 +1200,7 @@ class MoRIIOConnectorScheduler: self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], - block_notify_list=block_notify_list, + block_notify_list=blocks.get_block_ids()[0], host=_notify_host, port=target_port, ) @@ -1541,8 +1550,9 @@ class MoRIIOConnectorScheduler: * An ACK that arrives BEFORE its request finished is parked in ``_pending_sent_acks`` and released on a later step once the request enters ``_deferred_send_deadlines``. - * A deferred send whose ACK never arrives is reaped after - ``_defer_timeout`` and surfaced, so leaked blocks are freed. + * In READ mode, a deferred send whose ACK never arrives is reaped + after ``_defer_timeout``. WRITE-mode timeouts are owned by the worker, + which marks the transfer terminal before ACKing the block release. * A parked ACK that never matches a deferral before its own deadline is a stale duplicate (e.g. a real ACK landing after the send was already reaped) and is dropped. @@ -1575,7 +1585,7 @@ class MoRIIOConnectorScheduler: expired = [ req_id for req_id, (deadline, _) in self._deferred_send_deadlines.items() - if now >= deadline + if self.mode == MoRIIOMode.READ and now >= deadline ] if expired: safe.update(expired) @@ -1685,6 +1695,7 @@ class MoRIIOConnectorWorker: # Completions that arrived before transfer_id_to_request_id was populated. # Retried each step until the mapping is established. self._unmatched_write_completions: set[str] = set() + self._unmatched_write_failures: OrderedDict[TransferId, None] = OrderedDict() # Producer-side READ-mode ACK fan-in. When decode TP is larger than # prefill TP, multiple decode ranks can read from one prefill rank and # notify the same transfer_id. Blocks are reusable only after all ACKs. @@ -1877,6 +1888,7 @@ class MoRIIOConnectorWorker: remote_ip: str, multi_pod_hosts: list[str], remote_dp_size_local: int, + remote_dp_size: int, ) -> None: """Schedule a block write operation. @@ -1892,6 +1904,7 @@ class MoRIIOConnectorWorker: remote_ip: IP address of remote node multi_pod_hosts: List of pod IPs for multi-pod Wide-EP remote_dp_size_local: Per-pod DP size for multi-pod + remote_dp_size: Global DP size of the decode instance """ # synchronization to prevent dirty reads between @@ -1915,6 +1928,7 @@ class MoRIIOConnectorWorker: remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts, remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) self._writer.schedule_write(task) @@ -2686,6 +2700,28 @@ class MoRIIOConnectorWorker: return done_sending, done_recving + def get_transfer_results(self) -> KVConnectorTransferResults: + done_sending, done_recving = self.get_finished() + failed_recving: set[ReqId] = set() + + if self.mode == MoRIIOMode.WRITE and not self.is_producer: + pending = self._unmatched_write_failures + mapping = self.transfer_id_to_request_id + for transfer_id in self.moriio_wrapper.pop_failed_write_req_ids(): + pending[transfer_id] = None + for transfer_id in [t for t in pending if t in mapping]: + del pending[transfer_id] + failed_recving.add(mapping[transfer_id]) + while len(pending) > _MAX_UNMATCHED_WRITE_FAILURES: + pending.popitem(last=False) + done_recving |= failed_recving + + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recving, + ) + def wait_for_layer_load(self, layer_name: str) -> None: """Block until all in-flight READs of this layer have landed. @@ -3258,6 +3294,7 @@ class MoRIIOConnectorWorker: remote_ip=meta.remote_host, multi_pod_hosts=hosts, remote_dp_size_local=dp_local, + remote_dp_size=int(meta.remote_dp_size), ) def merge_contiguous_blocks( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 4ec3f72fb4..e33a1524e1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -33,8 +33,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( TransferError, TransferId, WriteTask, + fold_local_rank, get_port_offset, get_role, + pod_index, zmq_ctx, ) @@ -226,15 +228,6 @@ class MoRIIOWriter: for task in self._deferred_tasks: if self._is_transfer_terminal(task.transfer_id): continue - if now - task.enqueue_time > defer_timeout: - logger.error( - "Deferred write task for request %s expired after %.1fs " - "(remote blocks never arrived), marking done", - task.request_id, - now - task.enqueue_time, - ) - self._mark_request_done(task.transfer_id) - continue if self._is_remote_ready(task): try: self._execute_write_task(task) @@ -244,11 +237,42 @@ class MoRIIOWriter: task.request_id, ) self._mark_request_done(task.transfer_id) + elif ( + now - task.enqueue_time > defer_timeout + and self.worker.moriio_wrapper.fail_unallocated_write(task.transfer_id) + ): + self._notify_write_failed(task, now - task.enqueue_time) else: still_deferred.append(task) self._deferred_tasks = still_deferred + def _notify_write_failed(self, task: WriteTask, age: float) -> None: + self._clear_transfer_state(task.transfer_id) + logger.error( + "Deferred write task for request %s timed out after %.1fs waiting " + "for remote blocks", + task.request_id, + age, + ) + # No remote allocation arrived, so the decode DP rank that owns the + # request is unknown. Tell every rank; the others drop it. + for decode_dp_rank in range(max(task.remote_dp_size, 1)): + remote_ip, remote_port = self._resolve_notify_endpoint(task, decode_dp_rank) + try: + self.worker.moriio_wrapper.send_notify( + task.transfer_id, + remote_ip, + remote_port, + message_type="write_failed", + ) + except Exception: + logger.exception( + "Failed to notify decode DP rank %d that transfer %s timed out", + decode_dp_rank, + task.transfer_id, + ) + def _clear_transfer_state(self, transfer_id: TransferId) -> None: with self._write_state_lock: self._scheduled_writes.pop(transfer_id, None) @@ -314,27 +338,12 @@ class MoRIIOWriter: request_info = self._get_remote_alloc_info(task.transfer_id) with self._write_state_lock: request_info.completion_request_id = task.request_id - # Wide-EP multi-pod: task.remote_ip addresses only the first pod, - # so resolve the per-rank host from multi_pod_hosts (falls back to - # task.remote_ip for single-pod or on any indexing miss). - _remote_ip = task.remote_ip - _hosts = task.multi_pod_hosts - _dp_local = task.remote_dp_size_local - if _hosts and _dp_local > 0: - _pod_idx = int(request_info.decode_dp_rank) // _dp_local - if 0 <= _pod_idx < len(_hosts): - _remote_ip = _hosts[_pod_idx] - request_info.completion_remote_ip = _remote_ip - # Resolve the final notify port under the lock so finalize only - # reads it. The offset uses the per-pod local rank (% dp_local), - # since each pod binds notify sockets only for its local ranks; - # single-pod is bit-identical (modulus is a no-op). - _decode_dp_rank_for_port = int(request_info.decode_dp_rank) - if _dp_local > 0: - _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local - request_info.completion_notify_port = task.remote_notify_port + ( - get_port_offset(_decode_dp_rank_for_port, self.worker.tp_rank) - ) + # Resolve the final notify endpoint under the lock so finalize only + # reads it. + ( + request_info.completion_remote_ip, + request_info.completion_notify_port, + ) = self._resolve_notify_endpoint(task, request_info.decode_dp_rank) if task.transfer_id in self._sealed_writes: request_info.writes_expected = self._sealed_writes[task.transfer_id] @@ -506,6 +515,26 @@ class MoRIIOWriter: remote_port, ) + def _resolve_notify_endpoint( + self, task: WriteTask, decode_dp_rank: int + ) -> tuple[str, int]: + # Wide-EP multi-pod: task.remote_ip addresses only the first pod, so + # resolve the per-rank host from multi_pod_hosts (falls back to + # task.remote_ip for single-pod or on any indexing miss). The port + # offset uses the per-pod local rank, since each pod binds notify + # sockets only for its local ranks. + decode_dp_rank = int(decode_dp_rank) + dp_local = task.remote_dp_size_local + remote_ip = task.remote_ip + if task.multi_pod_hosts and dp_local > 0: + remote_pod = pod_index(decode_dp_rank, dp_local) + if 0 <= remote_pod < len(task.multi_pod_hosts): + remote_ip = task.multi_pod_hosts[remote_pod] + remote_port = task.remote_notify_port + get_port_offset( + fold_local_rank(decode_dp_rank, dp_local), self.worker.tp_rank + ) + return remote_ip, remote_port + class MoRIIOWrapper: """Wrapper for MoRIIO engine operations. @@ -539,6 +568,7 @@ class MoRIIOWrapper: self.done_req_ids: list[MoRIIOTransferAck] = [] self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {} self.done_write_cache_req_ids: list[str] = [] + self.failed_write_cache_req_ids: list[str] = [] self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict() self._transfer_timeout = transfer_timeout self.notify_thread: threading.Thread | None = None @@ -803,6 +833,8 @@ class MoRIIOWrapper: self._handle_remote_blocks_message(data) elif message_type == "write_done": self._handle_write_done_message(data) + elif message_type == "write_failed": + self._handle_write_failed_message(data) elif message_type == "release": self._handle_release_message(data) else: @@ -837,6 +869,12 @@ class MoRIIOWrapper: with self.lock: self.done_write_cache_req_ids.append(transfer_id) + def _handle_write_failed_message(self, data: dict): + assert get_role() != ROLE.PRODUCER, "Only decode can get WRITE failure messages" + transfer_id = data["transfer_id"] + with self.lock: + self.failed_write_cache_req_ids.append(transfer_id) + def _handle_release_message(self, data: dict): assert get_role() == ROLE.PRODUCER, ( "Only prefill can get transfer release messages" @@ -852,6 +890,23 @@ class MoRIIOWrapper: self.done_remote_allocate_req_dict.pop(transfer_id, None) self._mark_transfer_terminal_locked(transfer_id) + def fail_unallocated_write(self, transfer_id: TransferId) -> bool: + """Mark a WRITE transfer terminal and release its source blocks if + decode never sent its allocation. + + Returns False if the transfer is already terminal or the allocation + has arrived; the writer then drops or executes the task as usual. + """ + with self.lock: + if ( + self._is_transfer_terminal_locked(transfer_id) + or transfer_id in self.done_remote_allocate_req_dict + ): + return False + self._mark_transfer_terminal_locked(transfer_id) + self.done_req_ids.append(MoRIIOTransferAck(transfer_id)) + return True + def _handle_completion_message(self, msg: str): with self.lock: if get_role() == ROLE.PRODUCER: @@ -929,6 +984,12 @@ class MoRIIOWrapper: self.done_write_cache_req_ids = [] return done_write_cache + def pop_failed_write_req_ids(self): + with self.lock: + failed_write_cache = set(self.failed_write_cache_req_ids) + self.failed_write_cache_req_ids = [] + return failed_write_cache + def shutdown(self): logger.debug("Closing MoRIIOWrapper and cleaning up ZMQ sockets") for path, sock in self.paths.items():