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:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user