diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 41ed82a87..68b6af232 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) ):