Extract collect_inflight_reqs from abort_request for reusing (#35923)

This commit is contained in:
fzyzcjy
2026-08-24 20:15:09 +08:00
committed by GitHub
parent 9856b58de4
commit a37fdae562
+10 -7
View File
@@ -4570,6 +4570,15 @@ class Scheduler(
barrier(group=self.tp_group.cpu_group)
return RpcReqOutput(success=success, message="" if not exec else str(exec))
def collect_inflight_reqs(self) -> Set[Req]:
if self.ps.pp_size == 1:
inflight_batches = [self.running_batch, self.last_batch]
else:
inflight_batches = [*self.running_mbs, *self.mbs]
return {
req for batch in inflight_batches if batch is not None for req in batch.reqs
}
def abort_request(self, recv_req: AbortReq):
if (chunked_req := self.chunked_req) is not None:
if recv_req.abort_all or chunked_req.rid.startswith(recv_req.rid):
@@ -4711,13 +4720,7 @@ class Scheduler(
self.disagg_decode_prealloc_queue.retracted_queue = remaining_retracted
# Delete requests in the running batch
if self.ps.pp_size == 1:
inflight_batches = [self.running_batch, self.last_batch]
else:
inflight_batches = [*self.running_mbs, *self.mbs]
inflight_reqs = {r for b in inflight_batches if b is not None for r in b.reqs}
for req in inflight_reqs:
for req in self.collect_inflight_reqs():
if not req.finished() and (
recv_req.abort_all or req.rid.startswith(recv_req.rid)
):