Extract collect_inflight_reqs from abort_request for reusing (#35923)
This commit is contained in:
@@ -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)
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user