From a37fdae56213e7513389ba51187125a6adb62133 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:15:09 +0800 Subject: [PATCH] Extract collect_inflight_reqs from abort_request for reusing (#35923) --- python/sglang/srt/managers/scheduler.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) 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) ):