[PD] Clean early abort logic in PD module (#25677)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
@@ -607,7 +607,7 @@ class DecodePreallocQueue:
|
||||
|
||||
# Still poll if any receiver was aborted, otherwise it stays stuck.
|
||||
if all(decode_req.waiting_for_input for decode_req in self.queue) and not any(
|
||||
getattr(decode_req.kv_receiver, "conclude_state", None) == KVPoll.Failed
|
||||
decode_req.kv_receiver.conclude_state == KVPoll.Failed
|
||||
for decode_req in self.queue
|
||||
):
|
||||
return
|
||||
|
||||
@@ -3432,16 +3432,12 @@ class Scheduler(
|
||||
if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
|
||||
logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
|
||||
decode_req.kv_receiver.abort()
|
||||
if not isinstance(decode_req.req.finished_reason, FINISH_ABORT):
|
||||
decode_req.req.finished_reason = FINISH_ABORT()
|
||||
|
||||
# Abort requests waiting for kvcache to release tree cache
|
||||
for decode_req in self.disagg_decode_transfer_queue.queue:
|
||||
if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
|
||||
logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}")
|
||||
decode_req.kv_receiver.abort()
|
||||
if not isinstance(decode_req.req.finished_reason, FINISH_ABORT):
|
||||
decode_req.req.finished_reason = FINISH_ABORT()
|
||||
|
||||
# Abort requests already retracted to CPU cache
|
||||
if self.disagg_decode_prealloc_queue.retracted_queue:
|
||||
|
||||
Reference in New Issue
Block a user