[Bugfix][Mooncake] Report failed remote KV loads to the scheduler (#50984)
Signed-off-by: liuqihao <liuqihao@myhexin.com> Co-authored-by: liuqihao <liuqihao@myhexin.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
BruceLoveDecimal committed
Sep 14, 2026 at 17:55 UTC
2c7ee87223f0c94f614a0b8600f6f193e6e4388f
1 file changed
+66
-14
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+66
-14
@@ -2,9 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import asyncio
4
import logging
5
+import queue
6
import threading
7
import time
8
from collections import defaultdict
9
+from collections.abc import Collection
10
from concurrent.futures import ThreadPoolExecutor
11
from dataclasses import dataclass
12
from enum import IntEnum
@@ -433,6 +435,9 @@ class PullReqMeta:
435
expire_time: float = float("inf")
436
# Designed for one D pairing to multiple P
437
pull_tasks_count: int = 0
438
+ # Set once any worker reports a failure, so a success from another worker
439
+ # for the same request is not counted afterwards.
440
+ failed: bool = False
441
442
443
@dataclass
@@ -582,6 +587,11 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
587
assert self.connector_worker is not None
588
return self.connector_worker.get_finished()
589
590
+ def get_block_ids_with_load_errors(self) -> set[int]:
591
+ """Get block IDs whose remote KV load failed."""
592
+ assert self.connector_worker is not None
593
+ return self.connector_worker.get_block_ids_with_load_errors()
594
+
595
def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
596
assert self.connector_worker is not None
597
assert isinstance(self._connector_metadata, MooncakeConnectorMetadata)
@@ -1034,6 +1044,8 @@ class MooncakeConnectorWorker:
1044
1045
self.finished_sending_reqs: set[ReqId] = set()
1046
self.finished_recving_reqs: set[ReqId] = set()
1047
+ # Written from the receiver loop, drained from the worker thread.
1048
+ self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
1049
1050
self.xfer_stats = MooncakeKVConnectorStats()
1051
@@ -1907,12 +1919,11 @@ class MooncakeConnectorWorker:
1919
ret_msg = await sock.recv()
1920
response = self._xfer_resp_decoder.decode(ret_msg)
1921
if response.status == MooncakeXferResponseStatus.ERROR:
1910
- logger.error(
1911
- "Error happens during transferring kvcache for %s: %s",
1922
+ self._handle_failed_recv(
1923
+ pull_metas,
1924
req_ids,
1913
- response.err_msg,
1925
+ response.err_msg or "transfer error",
1926
)
1915
- self.xfer_stats.record_failed_recv()
1927
return
1928
self.process_pulling_result(response, pull_metas)
1929
if response.status == MooncakeXferResponseStatus.FINISH:
@@ -1920,10 +1931,50 @@ class MooncakeConnectorWorker:
1931
except zmq.ContextTerminated:
1932
logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
1933
except Exception as e:
1923
- logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
1924
- self.xfer_stats.record_failed_recv()
1934
+ self._handle_failed_recv(pull_metas, req_ids, f"transfer failed: {e}")
1935
return
1936
1937
+ def _handle_failed_recv(
1938
+ self,
1939
+ pull_metas: dict[ReqId, PullReqMeta],
1940
+ req_ids: Collection[ReqId],
1941
+ reason: str,
1942
+ ) -> None:
1943
+ """Report a failed remote KV load so the scheduler can fail or recompute it."""
1944
+ failed: list[ReqId] = []
1945
+ for req_id in req_ids:
1946
+ pull_meta = pull_metas.get(req_id)
1947
+ if pull_meta is None or pull_meta.failed:
1948
+ continue
1949
+ pull_meta.failed = True
1950
+ failed.append(req_id)
1951
+ self.xfer_stats.record_failed_recv()
1952
+
1953
+ invalid = {b for group in pull_meta.local_block_ids for b in group}
1954
+ if not invalid:
1955
+ # A pull with no local blocks only asks P to release its blocks
1956
+ # for a request that never reached the scheduler (see
1957
+ # AsyncLLM.notify_kv_transfer_request_rejected, which submits an
1958
+ # abort_immediately request just to run request_finished). No D
1959
+ # request is waiting on a load, and reporting one here would trip
1960
+ # the scheduler's `assert req_id in self.requests`.
1961
+ continue
1962
+ self._invalid_block_ids.put(invalid)
1963
+ self.finished_recving_reqs.add(pull_meta.d_req_id)
1964
+
1965
+ if failed:
1966
+ logger.error("pulling kv_caches for %s failed: %s", failed, reason)
1967
+
1968
+ def get_block_ids_with_load_errors(self) -> set[int]:
1969
+ """Drain the blocks whose remote KV load failed since the last call."""
1970
+ result: set[int] = set()
1971
+ while True:
1972
+ try:
1973
+ result.update(self._invalid_block_ids.get_nowait())
1974
+ except queue.Empty:
1975
+ break
1976
+ return result
1977
+
1978
def process_pulling_result(
1979
self,
1980
response: MooncakeXferResponse,
@@ -1933,6 +1984,8 @@ class MooncakeConnectorWorker:
1984
1985
for req_id in ok_reqs:
1986
pull_meta = pull_metas[req_id]
1987
+ if pull_meta.failed:
1988
+ continue
1989
# No race because we are in async loop.
1990
pull_meta.pull_tasks_count -= 1
1991
if pull_meta.pull_tasks_count == 0:
@@ -1942,10 +1995,8 @@ class MooncakeConnectorWorker:
1995
logger.debug("pulling kv_caches for %s finished", ok_reqs)
1996
1997
if response.err_reqs:
1945
- logger.error(
1946
- "pulling kv_caches for %s failed: %s",
1947
- response.err_reqs,
1948
- response.err_msg,
1998
+ self._handle_failed_recv(
1999
+ pull_metas, response.err_reqs, response.err_msg or "unknown error"
2000
)
2001
2002
async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
@@ -2023,10 +2074,11 @@ class MooncakeConnectorWorker:
2074
await self._pending_bootstrap_queries[remote_bootstrap_addr].wait()
2075
2076
if remote_engine_id not in self._remote_agents:
2026
- logger.error(
2027
- "Failed to find remote engine_id %s from bootstrap server %s",
2028
- remote_engine_id,
2029
- remote_bootstrap_addr,
2077
+ self._handle_failed_recv(
2078
+ pull_metas,
2079
+ list(pull_metas),
2080
+ f"remote engine_id {remote_engine_id} not found from bootstrap "
2081
+ f"server {remote_bootstrap_addr}",
2082
)
2083
return
2084