[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 <cursoragent@cursor.com>
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>

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"]

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

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(

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():