[PD] Fix KV cache corruption on abort by notifying ongoing prefill (#27372)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
@@ -877,6 +877,7 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
self.conclude_state: Optional[KVPoll] = None
|
||||
self.require_staging: bool = False
|
||||
self.init_time: Optional[float] = None
|
||||
self.abort_notified: bool = False
|
||||
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
|
||||
|
||||
@@ -1062,6 +1063,13 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
f"in KVPoll.WaitingForInput",
|
||||
)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
if (
|
||||
not self.abort_notified
|
||||
and hasattr(self, "bootstrap_infos")
|
||||
and self.bootstrap_infos is not None
|
||||
):
|
||||
self._send_abort_notification()
|
||||
self.abort_notified = True
|
||||
return KVPoll.Failed
|
||||
|
||||
def failure_exception(self):
|
||||
@@ -1079,6 +1087,36 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
if (
|
||||
not self.abort_notified
|
||||
and hasattr(self, "bootstrap_infos")
|
||||
and self.bootstrap_infos is not None
|
||||
):
|
||||
self._send_abort_notification()
|
||||
self.abort_notified = True
|
||||
|
||||
def _send_abort_notification(self):
|
||||
for bootstrap_info in self.bootstrap_infos:
|
||||
# Best-effort notification to prefill side that this request was aborted.
|
||||
try:
|
||||
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
||||
with lock:
|
||||
sock.send_multipart(
|
||||
[
|
||||
b"ABORT",
|
||||
str(self.bootstrap_room).encode("ascii"),
|
||||
self.kv_mgr.local_ip.encode("ascii"),
|
||||
str(self.kv_mgr.rank_port).encode("ascii"),
|
||||
]
|
||||
)
|
||||
logger.debug(
|
||||
f"Sent abort notification for room {self.bootstrap_room} "
|
||||
f"to {bootstrap_info.get('rank_ip', 'unknown')}:{bootstrap_info.get('rank_port', 'unknown')}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
f"Failed to send abort notification for room {self.bootstrap_room}: {e}"
|
||||
)
|
||||
|
||||
|
||||
class CommonKVBootstrapServer(BaseKVBootstrapServer):
|
||||
|
||||
@@ -1168,6 +1168,21 @@ class MooncakeKVManager(CommonKVManager):
|
||||
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
|
||||
)
|
||||
|
||||
if (
|
||||
kv_chunk.room not in self.request_status
|
||||
or self.check_status(kv_chunk.room) == KVPoll.Failed
|
||||
):
|
||||
logger.debug(
|
||||
f"Skipping chunk for room {kv_chunk.room} because it has already failed or been aborted"
|
||||
)
|
||||
if self.enable_trace:
|
||||
kv_chunk.trace_ctx.trace_slice_end(
|
||||
MooncakeRequestStage.MOONCAKE_WORKER_SEND.stage_name,
|
||||
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
|
||||
thread_finish_flag=True,
|
||||
)
|
||||
continue
|
||||
|
||||
if (
|
||||
self.enable_staging
|
||||
and staging_strategy is None
|
||||
@@ -1386,6 +1401,44 @@ class MooncakeKVManager(CommonKVManager):
|
||||
|
||||
handle_staging_rsp(waiting_req_bytes, self.transfer_infos)
|
||||
continue
|
||||
# Decode-side abort notification: mark room as failed and ACK
|
||||
if room == "ABORT":
|
||||
room_to_be_aborted = int(waiting_req_bytes[1].decode("ascii"))
|
||||
decode_ip = waiting_req_bytes[2].decode("ascii")
|
||||
decode_port = int(waiting_req_bytes[3].decode("ascii"))
|
||||
# No need to abort the room if it has already succeeded
|
||||
if (
|
||||
room_to_be_aborted in self.request_status
|
||||
and self.check_status(room_to_be_aborted) != KVPoll.Success
|
||||
):
|
||||
self.update_status(room_to_be_aborted, KVPoll.Failed)
|
||||
logger.debug(
|
||||
f"Received abort notification for room {room_to_be_aborted}, "
|
||||
f"marked as Failed"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"Received abort notification for room {room_to_be_aborted}, "
|
||||
f"ignoring (already completed or unknown)"
|
||||
)
|
||||
# Send ACK back to decode endpoint
|
||||
try:
|
||||
na = NetworkAddress(decode_ip, decode_port)
|
||||
self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart(
|
||||
[
|
||||
b"ABORT_ACK",
|
||||
str(room_to_be_aborted).encode("ascii"),
|
||||
]
|
||||
)
|
||||
logger.debug(
|
||||
f"Sent ABORT_ACK for room {room_to_be_aborted} to "
|
||||
f"{decode_ip}:{decode_port}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
f"Failed to send ABORT_ACK for room {room_to_be_aborted}: {e}"
|
||||
)
|
||||
continue
|
||||
mooncake_session_id = waiting_req_bytes[3].decode("ascii")
|
||||
if room == "None":
|
||||
self.decode_kv_args_table[mooncake_session_id] = (
|
||||
@@ -1457,6 +1510,13 @@ class MooncakeKVManager(CommonKVManager):
|
||||
self._handle_staging_req(msg)
|
||||
continue
|
||||
|
||||
# Prefill acknowledges abort notification
|
||||
if msg[0] == b"ABORT_ACK":
|
||||
# TODO(shangming): use this info to implement the deferred release mechanism if needed
|
||||
ack_aborted_room = int(msg[1].decode("ascii"))
|
||||
logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}")
|
||||
continue
|
||||
|
||||
bootstrap_room, status, prefill_rank = msg
|
||||
status = int(status.decode("ascii"))
|
||||
bootstrap_room = int(bootstrap_room.decode("ascii"))
|
||||
|
||||
Reference in New Issue
Block a user