fix(disagg/mooncake): mark abort Failed and skip aborted chunks in worker

Port the abort-correctness pair from upstream sgl-project/sglang:
- #24522: abort() now calls kv_mgr.update_status(room, Failed) so the
  transfer worker and other pollers observe the abort, not only the local
  conclude_state.
- #27372 (guard half): the transfer worker skips a chunk whose room has
  already failed or been aborted, so it never transfers into KV pages that
  may have been reclaimed. This matters more once per-layer async transfer
  widens the abort-vs-in-flight window.

The decode->prefill ABORT/ABORT_ACK notification half of #27372 is deferred:
it interleaves with staging/tracing infra this branch lacks and needs the PD
test harness to verify, so porting it now would be unverified guesswork.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-06 22:03:52 +00:00
parent 0cd5c38a9e
commit 2524fe4c9d

View File

@@ -1000,6 +1000,21 @@ class MooncakeKVManager(CommonKVManager):
while True:
try:
kv_chunk: TransferKVChunk = queue.get()
# Skip a chunk whose room has already failed or been aborted, so we never
# transfer into KV pages that may have been reclaimed. Port of upstream
# sgl-project/sglang #27372 (abort KV-cache corruption guard). The
# decode->prefill ABORT/ABORT_ACK notification half of that PR is deferred
# until the PD test harness exists (it interleaves with staging/tracing
# infra this branch does not carry).
if (
kv_chunk.room not in self.request_status
or self.check_status(kv_chunk.room) == KVPoll.Failed
):
logger.debug(
"Skipping chunk for room %s because it has already failed or been aborted",
kv_chunk.room,
)
continue
reqs_to_be_processed = (
self.transfer_infos[kv_chunk.room].values()
if kv_chunk.room in self.transfer_infos
@@ -1580,6 +1595,10 @@ class MooncakeKVSender(CommonKVSender):
self.bootstrap_room,
"Aborted by AbortReq.",
)
# Mark the manager status Failed (not only the local conclude_state) so the
# transfer worker and other pollers observe the abort. Port of upstream
# sgl-project/sglang #24522.
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
# Explicitly set the status to failure since this request has been aborted
self.conclude_state = KVPoll.Failed
@@ -1760,6 +1779,10 @@ class MooncakeKVReceiver(CommonKVReceiver):
self.bootstrap_room,
"Aborted by AbortReq.",
)
# Mark the manager status Failed (not only the local conclude_state) so the
# transfer worker and other pollers observe the abort. Port of upstream
# sgl-project/sglang #24522.
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
# Explicitly set the status to failure since this request has been aborted
self.conclude_state = KVPoll.Failed