[Refactor] scheduler: Share one interface across request lists and abort through each list

python/sglang/srt/constrained/grammar_manager.py

@@ -14,6 +14,10 @@
 from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
 from sglang.srt.distributed.communication_tags import P2PTag
 from sglang.srt.environ import envs
+from sglang.srt.managers.scheduler_components.request_list import (
+    ListKind,
+    _RequestList,
+)
 from sglang.srt.runtime_context import get_parallel, get_serving
 from sglang.srt.sampling.sampling_params import (
     get_request_reasoning_end_token_ids,
@@ -32,7 +36,7 @@ class GrammarManager:
     def __init__(self, scheduler: Scheduler):
         self.scheduler = scheduler
         self.server_args = scheduler.server_args
-        self.grammar_queue: List[Req] = []
+        self.grammar_queue = _RequestList(ListKind.GRAMMAR, abort=self._abort)
         if not get_serving().skip_tokenizer_init:
             self.grammar_backend = create_grammar_backend(
                 self.server_args,
@@ -111,9 +115,11 @@ def _pp_sync_ready_failed(
             )
         return data
 
-    def abort_requests(self, recv_req: AbortReq):
+    def _abort(self, recv_req: AbortReq) -> None:
+        # An aborted request still runs one prefill pass, which
+        # `set_finish_with_abort` cuts to one token.
         for req in self.grammar_queue:
-            if recv_req.abort_all or req.rid.startswith(recv_req.rid):
+            if recv_req.matches(req.rid):
                 logger.debug(f"Abort grammar queue request. {req.rid=}")
                 if isinstance(req.grammar, futures.Future) and req.grammar:
                     req.grammar.cancel()
@@ -204,7 +210,7 @@ def process_req_with_grammar(self, req: Req) -> bool:
                 self._apply_request_reasoning_config(req)
 
         if add_to_grammar_queue:
-            self.grammar_queue.append(req)
+            self.grammar_queue.join([req])
 
         return add_to_grammar_queue
 
@@ -224,14 +230,16 @@ def get_ready_grammar_requests(self) -> List[Req]:
         ranks receive and apply the propagated ready/failed decision.
         """
         assert self.grammar_backend
+        # Indices into this snapshot are synced across ranks.
+        grammar_queue = list(self.grammar_queue)
         ready_req_idxs: set[int] = set()
         failed_req_idxs: set[int] = set()
 
         if self.pp_rank == 0:
             # Poll for ready requests
             start_time = time.perf_counter()
             while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL:
-                for i, req in enumerate(self.grammar_queue):
+                for i, req in enumerate(grammar_queue):
                     if i in ready_req_idxs:
                         continue
 
@@ -245,20 +253,20 @@ def get_ready_grammar_requests(self) -> List[Req]:
                     if req.grammar.done():
                         ready_req_idxs.add(i)
 
-                if len(ready_req_idxs) == len(self.grammar_queue):
+                if len(ready_req_idxs) == len(grammar_queue):
                     break
 
                 # Sleep a bit to avoid busy waiting
                 time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10)
 
             # Check failed requests
-            for i, req in enumerate(self.grammar_queue):
+            for i, req in enumerate(grammar_queue):
                 if i not in ready_req_idxs:
                     # grammar_wait_ct is only updated on PP0; later PP ranks
                     # receive PP0's ready/failed decision through PP sync.
-                    self.grammar_queue[i].grammar_wait_ct += 1
+                    grammar_queue[i].grammar_wait_ct += 1
                     if (
-                        self.grammar_queue[i].grammar_wait_ct
+                        grammar_queue[i].grammar_wait_ct
                         >= self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS
                     ):
                         # Timeout after max poll iterations
@@ -296,7 +304,7 @@ def get_ready_grammar_requests(self) -> List[Req]:
         # Return ready requests
         return_reqs: List[Req] = []
         for i in synced_ready_req_idxs:
-            req = self.grammar_queue[i]
+            req = grammar_queue[i]
             return_reqs.append(req)
             if req.finished() or req.grammar is None:  # It is aborted by AbortReq
                 continue
@@ -318,7 +326,7 @@ def get_ready_grammar_requests(self) -> List[Req]:
 
         # Return failed requests
         for i in synced_failed_req_idxs:
-            req = self.grammar_queue[i]
+            req = grammar_queue[i]
             return_reqs.append(req)
 
             assert isinstance(req.grammar, futures.Future) and req.grammar_key
@@ -330,9 +338,5 @@ def get_ready_grammar_requests(self) -> List[Req]:
             req.set_finish_with_abort(error_msg)
 
         # Remove finished requests from grammar_queue
-        self.grammar_queue = [
-            req
-            for i, req in enumerate(self.grammar_queue)
-            if i not in synced_ready_req_idxs and i not in synced_failed_req_idxs
-        ]
+        self.grammar_queue.leave(return_reqs)
         return return_reqs

python/sglang/srt/disaggregation/decode.py

@@ -90,7 +90,15 @@
     kv_private_swa_tokens,
     kv_private_tokens,
 )
-from sglang.srt.managers.utils import GenerationBatchResult, is_war_barrier_enabled
+from sglang.srt.managers.scheduler_components.request_list import (
+    ListKind,
+    _RequestList,
+)
+from sglang.srt.managers.utils import (
+    GenerationBatchResult,
+    is_war_barrier_enabled,
+    make_abort_req,
+)
 from sglang.srt.mem_cache.allocation import ensure_mamba_capacity
 from sglang.srt.mem_cache.allocation_sizing import get_mamba_tracking_slots
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -488,7 +496,9 @@ def __init__(
         self.transfer_backend = transfer_backend
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
-        self.retracted_queue: List[Req] = []
+        self.retracted_queue = _RequestList(
+            ListKind.PD_RETRACTED, abort=self._abort_retracted
+        )
         self.pending_reqs: List[DecodeRequest] = []
         # In-flight authoritative room -> DP-rank lookups, consumed below.
         self._prefill_dp_rank_queries: Dict[
@@ -501,7 +511,9 @@ def __init__(
         # Retracted requests staged for rebootstrap while generation is paused.
         # Enqueued into ``self.queue`` only on ``continue_generation`` so the
         # prefix KV is recomputed under the post-retract (updated) weights.
-        self.held_rebootstrap_reqs: List[Req] = []
+        self.held_rebootstrap_reqs = _RequestList(
+            ListKind.PD_HELD, abort=self._abort_held
+        )
         self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
         if self.enable_staging and self.is_mla_backend:
             raise RuntimeError(
@@ -918,7 +930,7 @@ def add(
 
         if is_retracted:
             req.retraction_mb_id = None
-            self.retracted_queue.append(req)
+            self.retracted_queue.join([req])
         else:
             decode_req = self._create_receiver_and_enqueue(
                 req, is_rebootstrap=is_rebootstrap
@@ -1009,12 +1021,11 @@ def hold_rebootstrap(self, req: Req) -> None:
         enqueue it from ``enqueue_held_rebootstrap`` on resume, so its prefix KV
         is recomputed by the prefill worker under the updated weights.
         """
-        self.held_rebootstrap_reqs.append(req)
+        self.held_rebootstrap_reqs.join([req])
 
     def enqueue_held_rebootstrap(self) -> None:
         """Enqueue all staged rebootstrap requests when generation resumes."""
-        held = self.held_rebootstrap_reqs
-        self.held_rebootstrap_reqs = []
+        held = self.held_rebootstrap_reqs.take(lambda req: True)
         for req in held:
             self.add(req, is_rebootstrap=True)
 
@@ -1055,23 +1066,32 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool:
             return True
         return False
 
-    def take_aborted(self, recv_req: AbortReq) -> Tuple[List[Req], List[Req]]:
-        def aborted(req: Req) -> bool:
-            return recv_req.abort_all or req.rid.startswith(recv_req.rid)
+    def _abort_retracted(self, recv_req: AbortReq) -> None:
+        # Retracted requests hold a KV backup.
+        for req in self.retracted_queue.take(lambda req: recv_req.matches(req.rid)):
+            discard_kv_cache_backup(
+                req,
+                self.tree_cache,
+                get_disagg().disaggregation_decode_retraction_backup,
+            )
+            self.scheduler.send_abort_output(make_abort_req(req), req)
 
-        retracted = [req for req in self.retracted_queue if aborted(req)]
-        held = [req for req in self.held_rebootstrap_reqs if aborted(req)]
-        self.retracted_queue = [req for req in self.retracted_queue if not aborted(req)]
-        self.held_rebootstrap_reqs = [
-            req for req in self.held_rebootstrap_reqs if not aborted(req)
-        ]
-        return retracted, held
+    def _abort_held(self, recv_req: AbortReq) -> None:
+        # Held requests hold no KV.
+        for req in self.held_rebootstrap_reqs.take(
+            lambda req: recv_req.matches(req.rid)
+        ):
+            req.finish_outside_step(FINISH_ABORT())
+            self.scheduler.release_aborted_request(req)
+            self.scheduler.send_abort_output(
+                make_abort_req(req, finished_reason=recv_req.finished_reason), req
+            )
 
-    def request_lists(self) -> Dict[str, List[Req]]:
+    def request_lists(self) -> Dict[ListKind, List[Req]]:
         return {
-            "decode_prealloc": [decode_req.req for decode_req in self.queue],
-            "decode_retracted": self.retracted_queue,
-            "decode_held": self.held_rebootstrap_reqs,
+            ListKind.PD_PREALLOC: [decode_req.req for decode_req in self.queue],
+            ListKind.PD_RETRACTED: self.retracted_queue,
+            ListKind.PD_HELD: self.held_rebootstrap_reqs,
         }
 
     def extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
@@ -1104,7 +1124,6 @@ def resume_retracted_reqs(
 
         # allocate memory
         resumed_reqs = []
-        indices_to_remove = set()
         swa_allocatable_tokens = None
         if self._uses_swa_reservation():
             full_allocatable_tokens, swa_allocatable_tokens = (
@@ -1115,7 +1134,7 @@ def resume_retracted_reqs(
                 count_retracted=False
             )
 
-        for i, req in enumerate(self.retracted_queue):
+        for req in self.retracted_queue:
             if rids_to_check is not None and req.rid not in rids_to_check:
                 continue
 
@@ -1148,7 +1167,6 @@ def resume_retracted_reqs(
                     break
 
             resumed_reqs.append(req)
-            indices_to_remove.add(i)
             req.is_retracted = False
             self._pre_alloc(req)
             full_allocatable_tokens = self._allocatable_token_budgets(
@@ -1175,11 +1193,7 @@ def resume_retracted_reqs(
                 get_disagg().disaggregation_decode_retraction_backup,
             )
 
-        self.retracted_queue = [
-            entry
-            for i, entry in enumerate(self.retracted_queue)
-            if i not in indices_to_remove
-        ]
+        self.retracted_queue.leave(resumed_reqs)
 
         return resumed_reqs
 
@@ -2851,11 +2865,11 @@ def __init__(
         # (decode_req, start_time, deadline, metadata_idx, required_acks).
         self._deferred_releases: List[Tuple[DecodeRequest, float, float, int, int]] = []
 
-    def request_lists(self) -> Dict[str, List[Req]]:
+    def request_lists(self) -> Dict[ListKind, List[Req]]:
         return {
-            "decode_transfer": [decode_req.req for decode_req in self.queue],
+            ListKind.PD_RECEIVING: [decode_req.req for decode_req in self.queue],
             # Aborted, with KV held until the transfer drains.
-            "decode_releasing": [
+            ListKind.PD_RELEASING: [
                 decode_req.req for decode_req, *_ in self._deferred_releases
             ],
         }

python/sglang/srt/disaggregation/encoder/receiver.py

@@ -2042,7 +2042,7 @@ def abort_waiting_requests(self, recv_req) -> None:
         existing FAIL channel. AbortReq is broadcast, so every TP rank does
         this and the status all-reduce stays consistent."""
         for waiting_req in self.waiting_list:
-            if not (recv_req.abort_all or waiting_req.rid.startswith(recv_req.rid)):
+            if not recv_req.matches(waiting_req.rid):
                 continue
             if waiting_req.status in (
                 WaitingMMRequestStatus.PENDING,