From 6c269807b40c5c37921f1e98ca983432605a5b2c Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:10:25 -0800 Subject: [PATCH 01/12] move _make_abort_req to managers utils as make_abort_req --- python/sglang/srt/managers/scheduler.py | 45 ++++++------------- python/sglang/srt/managers/utils.py | 23 +++++++++- ...test_priority_scheduling_disaggregation.py | 8 ++-- .../unit/managers/test_scheduler_timeouts.py | 4 +- .../unit/utils/test_weight_versions.py | 9 ++-- 5 files changed, 47 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 626ae741e86d..4d68d370c03d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -303,6 +303,7 @@ allocate_distinct_stream, is_health_check_generate_req, is_war_barrier_enabled, + make_abort_req, validate_input_length, ) from sglang.srt.mem_cache import kv_cache_builder @@ -394,10 +395,7 @@ ) from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter -from sglang.srt.utils.weight_versions import ( - compute_weight_version_spans, - record_weight_version_events, -) +from sglang.srt.utils.weight_versions import record_weight_version_events from sglang.utils import TypeBasedDispatcher, get_exception_traceback if is_mps(): @@ -2786,7 +2784,7 @@ def init_deferred_responses(self) -> None: disaggregation_mode=self.disaggregation_mode, output_streamer=self.output_streamer, send_abort=lambda req, finished_reason: self.send_abort_output( - _make_abort_req(req, finished_reason=finished_reason), req + make_abort_req(req, finished_reason=finished_reason), req ), ) @@ -3289,7 +3287,7 @@ def _handle_generate_request( return if error_msg := self.validate_dllm_request(req): - abort_req = _make_abort_req( + abort_req = make_abort_req( req, finished_reason={ "type": "abort", @@ -3694,7 +3692,7 @@ def _set_or_validate_priority(self, req: Req) -> bool: and req.priority is not None and self.abort_on_priority_when_disabled ): - abort_req = _make_abort_req( + abort_req = make_abort_req( req, finished_reason={ "type": "abort", @@ -3726,7 +3724,7 @@ def _abort_uocr_doc_windows( for req in reqs: self._release_aborted_request(req) self.ipc_channels.send_to_tokenizer.send_output( - _make_abort_req(req, finished_reason=finished_reason), req + make_abort_req(req, finished_reason=finished_reason), req ) def _abort_on_queued_limit(self, recv_req: Req) -> bool: @@ -3766,7 +3764,7 @@ def _abort_on_queued_limit(self, recv_req: Req) -> bool: ) self.send_abort_output( - _make_abort_req( + make_abort_req( req_to_abort, finished_reason={ "type": "abort", @@ -4010,7 +4008,7 @@ def process_pending_chunked_abort(self) -> None: # A client abort leaves the message to the tokenizer. finished_reason = req.discard_output_reason and reason.to_json() self.send_abort_output( - _make_abort_req(req, finished_reason=finished_reason), req + make_abort_req(req, finished_reason=finished_reason), req ) logger.debug(f"Abort chunked prefill request. {req.rid=}") @@ -4816,7 +4814,7 @@ def update_running_batch(self, batch: ScheduleBatch) -> Optional[ScheduleBatch]: for req in reqs_to_abort: abort_reason: FINISH_ABORT = req.to_finish self.send_abort_output( - _make_abort_req(req, finished_reason=abort_reason.to_json()), + make_abort_req(req, finished_reason=abort_reason.to_json()), req, ) for req in reqs_to_abort: @@ -6045,7 +6043,7 @@ def abort_request(self, recv_req: AbortReq): # Without the initiator's reason the tokenizer falls back to a # generic abort message. self.send_abort_output( - _make_abort_req(req, finished_reason=recv_req.finished_reason), req + make_abort_req(req, finished_reason=recv_req.finished_reason), req ) # For disaggregation decode mode, the request in the waiting queue has KV cache allocated. if self.disaggregation_mode == DisaggregationMode.DECODE: @@ -6073,7 +6071,7 @@ def abort_request(self, recv_req: AbortReq): if req.finished_reason is None: req.finish_outside_step(FINISH_ABORT()) self._release_aborted_request(req) - self.send_abort_output(_make_abort_req(req), req) + self.send_abort_output(make_abort_req(req), req) if req.kv.holds_kv or req.kv.holds_mamba: release_kv_cache(req, self.tree_cache, checkpoint=False) logger.debug(f"Abort dLLM queued request. {req.rid=}") @@ -6134,12 +6132,12 @@ def abort_request(self, recv_req: AbortReq): self.tree_cache, get_disagg().disaggregation_decode_retraction_backup, ) - self.send_abort_output(_make_abort_req(req), req) + self.send_abort_output(make_abort_req(req), req) for req in held: req.finish_outside_step(FINISH_ABORT()) self._release_aborted_request(req) self.send_abort_output( - _make_abort_req(req, finished_reason=recv_req.finished_reason), req + make_abort_req(req, finished_reason=recv_req.finished_reason), req ) # Delete requests in the running batch @@ -6258,7 +6256,7 @@ def pause_generation(self, recv_req: PauseGenerationReqInput): status_code=HTTPStatus.SERVICE_UNAVAILABLE, ) self.send_abort_output( - _make_abort_req(req, finished_reason=abort_reason.to_json()), req + make_abort_req(req, finished_reason=abort_reason.to_json()), req ) elif self.disaggregation_mode == DisaggregationMode.DECODE: req.time_stats.set_retract_time() @@ -6756,18 +6754,3 @@ def run_scheduler_process( scheduler.release_host_resources() # Last: anything above may still need a working communicator. abort_distributed_environment() - - -def _make_abort_req( - req: Req, finished_reason: Optional[FinishReasonDict] = None -) -> AbortReq: - return AbortReq( - rid=req.rid, - finished_reason=finished_reason, - weight_versions=compute_weight_version_spans( - req.weight_version_events, - current_version=get_serving().weight_version, - num_output_tokens=len(req.output_ids), - ), - prefill_weight_versions=req.prefill_weight_versions, - ) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index a9c592f3f279..2587ad47fad3 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -18,14 +18,20 @@ ) from sglang.srt.managers import io_struct from sglang.srt.managers.auxiliary_output import append_auxiliary_output +from sglang.srt.managers.io_struct import AbortReq, FinishReasonDict from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord from sglang.srt.model_executor.forward_batch_info import PPProxyTensors -from sglang.srt.runtime_context import get_spec, max_speculative_num_draft_tokens +from sglang.srt.runtime_context import ( + get_serving, + get_spec, + max_speculative_num_draft_tokens, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.state_capturer.base import TopkCaptureOutput from sglang.srt.utils.common import async_d2h as _async_d2h from sglang.srt.utils.common import is_cuda +from sglang.srt.utils.weight_versions import compute_weight_version_spans if TYPE_CHECKING: from sglang.srt.managers.async_d2h_copy_worker import HostCopyDone @@ -515,3 +521,18 @@ def compute_num_reserved_tokens() -> int: spec.speculative_eagle_topk * spec.speculative_num_steps, max_speculative_num_draft_tokens(), ) + + +def make_abort_req( + req: Req, finished_reason: Optional[FinishReasonDict] = None +) -> AbortReq: + return AbortReq( + rid=req.rid, + finished_reason=finished_reason, + weight_versions=compute_weight_version_spans( + req.weight_version_events, + current_version=get_serving().weight_version, + num_output_tokens=len(req.output_ids), + ), + prefill_weight_versions=req.prefill_weight_versions, + ) diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index cd5b01c38b5b..84b1abfe4466 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -130,7 +130,7 @@ def test_priority_disabled_abort_validation_applies_to_decode_mode(self): req = self._new_req(priority=10) with patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v0"), ): scheduler._add_request_to_queue(req) @@ -150,7 +150,7 @@ def test_priority_rejection_releases_early_prefetch(self): req = self._new_req(priority=10) with patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v0"), ): scheduler._add_request_to_queue(req) @@ -178,7 +178,7 @@ def test_full_queue_releases_only_the_rejected_request(self): incoming = self._new_req(priority=priority) with patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v0"), ): scheduler._add_request_to_queue(incoming) @@ -458,7 +458,7 @@ def test_scheduler_waiting_abort_delegates_release_once(self): scheduler.tree_cache = MagicMock() with ( - patch("sglang.srt.managers.scheduler.get_serving") as get_serving, + patch("sglang.srt.managers.utils.get_serving") as get_serving, patch("sglang.srt.managers.scheduler.release_kv_cache") as release, ): get_serving.return_value.weight_version = None diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 94564c1acdf4..237450b97136 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -100,7 +100,7 @@ def _scheduler(waiting_queue, running_reqs=(), last_batch_reqs=()): class TestQueuedLimitAbort(CustomTestCase): def setUp(self): patcher = patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v0"), ) patcher.start() @@ -158,7 +158,7 @@ class TestQueuedAbort(CustomTestCase): def setUp(self): enter_scope(self, published_topology()) patcher = patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v0"), ) patcher.start() diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index c4cab051c7b8..315b167be8bf 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -6,10 +6,11 @@ from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.io_struct import AbortReq -from sglang.srt.managers.scheduler import Scheduler, _make_abort_req +from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) +from sglang.srt.managers.utils import make_abort_req from sglang.srt.utils.weight_versions import ( WeightVersionEvent, WeightVersionSpan, @@ -493,7 +494,7 @@ def test_priority_disabled_rejection_attaches_spans(self): ) with patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v1"), ): accepted = Scheduler._set_or_validate_priority(scheduler, req) @@ -513,10 +514,10 @@ def test_abort_req_carries_spans_clamped_to_sampled_tokens(self): WeightVersionEvent(old_version="v1", num_output_tokens=3) ) with patch( - "sglang.srt.managers.scheduler.get_serving", + "sglang.srt.managers.utils.get_serving", return_value=SimpleNamespace(weight_version="v2"), ): - abort_req = _make_abort_req(req) + abort_req = make_abort_req(req) self.assertIsInstance(abort_req, AbortReq) self.assertEqual(abort_req.rid, "r0") From fefa845aecf15cc9e1bc3e0f22071146ebd0618a Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:10:46 -0800 Subject: [PATCH 02/12] add RequestList trait, ListKind and move; key request lists by kind --- python/sglang/srt/disaggregation/decode.py | 15 ++-- .../srt/managers/hisparse_coordinator.py | 5 +- python/sglang/srt/managers/scheduler.py | 60 ++++++++----- .../deferred_responses.py | 5 +- .../scheduler_components/invariant_checker.py | 15 ++-- .../scheduler_components/request_list.py | 90 ++++++++++++++++++- .../scheduler_runtime.py | 5 +- .../test_disaggregation_optimistic_prefill.py | 7 +- .../test_prefill_abort_result_cleanup.py | 7 +- .../test_prefill_complete_integration.py | 7 +- .../test_prefill_result_polling.py | 13 +-- test/registered/unit/lora/test_lora_pp.py | 7 +- .../test_invariant_checker.py | 7 +- .../unit/managers/test_deferred_output.py | 7 +- .../test_disagg_idle_step_counters.py | 9 +- .../unit/managers/test_kv_page_invariants.py | 17 ++-- .../test_scheduler_chunked_req_gate.py | 7 +- ...est_scheduler_flush_cache_after_retract.py | 7 +- .../test_scheduler_held_row_admission.py | 7 +- .../managers/test_scheduler_hicache_events.py | 7 +- .../managers/test_scheduler_on_idle_load.py | 7 +- .../test_scheduler_pause_generation.py | 9 +- .../managers/test_scheduler_running_reqs.py | 15 ++-- .../test_scheduler_weight_version_tracking.py | 9 +- 24 files changed, 252 insertions(+), 92 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a0a207981cac..4e10a64ff80f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -90,6 +90,7 @@ kv_private_swa_tokens, kv_private_tokens, ) +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.managers.utils import GenerationBatchResult, is_war_barrier_enabled from sglang.srt.mem_cache.allocation import ensure_mamba_capacity from sglang.srt.mem_cache.allocation_sizing import get_mamba_tracking_slots @@ -1067,11 +1068,11 @@ def aborted(req: Req) -> bool: ] return retracted, held - 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: @@ -2851,11 +2852,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 ], } diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 01c5096d6bf0..8f7459279dcb 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -39,6 +39,7 @@ def load_blocks_to_device_buffer_mha(*args, **kwargs): from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator, @@ -1057,8 +1058,8 @@ def _grow_device_buffers( reserved_positions = (seq_lens - 1).clamp(max=self.device_buffer_size) return self.req_to_device_buffer[req_pool_indices, reserved_positions] - def request_lists(self) -> Dict[str, List[Req]]: - return {"hisparse_staging": [act.req for act in self.ack_staging_queue]} + def request_lists(self) -> Dict[ListKind, List[Req]]: + return {ListKind.HISPARSE_STAGING: [act.req for act in self.ack_staging_queue]} def has_ongoing_staging(self) -> bool: return len(self.ack_staging_queue) > 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4d68d370c03d..eed16311f4b3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -25,7 +25,18 @@ from contextlib import ExitStack, contextmanager, nullcontext from functools import partial from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Deque, + Dict, + Iterable, + List, + Optional, + Set, + Tuple, + Union, +) from sglang.srt.runtime_context import ( SpawnRanks, @@ -288,7 +299,10 @@ from sglang.srt.managers.scheduler_components.recv_skipper import ( SchedulerRecvSkipper, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) @@ -1369,11 +1383,11 @@ def init_running_status(self): self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False) # Requests merged into a running batch, until they finish or a retract, # preemption or pause takes them out. - self.running_reqs = RequestList() + self.running_reqs = _RequestList(ListKind.DECODING) # Prefill requests out of the chunked slot whose launched steps their next # holder has not taken yet: the running batch at the merge, the PD # transfer queue, HiSparse staging, or a requeue. - self.prefill_launched_reqs = RequestList() + self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) # The current forward batch self.cur_batch_for_debug: Optional[ScheduleBatch] = None # The last forward batch @@ -4666,24 +4680,24 @@ def _get_all_pp_reqs(self) -> List[Req]: if not req.finished() ] - def request_lists(self) -> Dict[str, List[Req]]: - """The lists that hold the live requests, by name. A request is in at + def request_lists(self) -> Dict[ListKind, Iterable[Req]]: + """The lists that hold the live requests, by kind. A request is in at most one.""" lists = { - "waiting": self.waiting_queue, - "grammar": self.grammar_manager.grammar_queue, - "chunked": [] if self.chunked_req is None else [self.chunked_req], - "prefill_launched": list(self.prefill_launched_reqs), - "running": list(self.running_reqs), + ListKind.WAITING: self.waiting_queue, + ListKind.GRAMMAR: self.grammar_manager.grammar_queue, + ListKind.CHUNKED: [] if self.chunked_req is None else [self.chunked_req], + ListKind.PD_PREFILLING: self.prefill_launched_reqs, + ListKind.DECODING: self.running_reqs, } # Read at call time: a PD role switch rebuilds the PD queues. if self.disagg_prefill_bootstrap_queue is not None: - lists["prefill_bootstrap"] = self.disagg_prefill_bootstrap_queue.queue + lists[ListKind.PD_BOOTSTRAP] = self.disagg_prefill_bootstrap_queue.queue if self.disagg_prefill_transfer_queue is not None: - lists["prefill_transfer"] = self.disagg_prefill_transfer_queue + lists[ListKind.PD_SENDING] = self.disagg_prefill_transfer_queue if self.dllm_config is not None: # Its staging queue is the part of this queue the current round runs. - lists["dllm"] = self.dllm_manager.waiting_queue + lists[ListKind.DLLM] = self.dllm_manager.waiting_queue components = ( self.disagg_decode_prealloc_queue, self.disagg_decode_transfer_queue, @@ -5669,25 +5683,25 @@ def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool: lists = self.request_lists() if for_health_check: - busy_lists = {"chunked", "waiting"} + busy_lists = {ListKind.CHUNKED, ListKind.WAITING} if not self._engine_paused: - busy_lists.add("decode_retracted") + busy_lists.add(ListKind.PD_RETRACTED) else: busy_lists = lists.keys() - { # Their steps are counted through the batches above. - "running", - "prefill_launched", + ListKind.DECODING, + ListKind.PD_PREFILLING, # Only its staging part is counted, above. - "dllm", + ListKind.DLLM, # Held across a weight update, whose cache flush needs idle. - "decode_held", + ListKind.PD_HELD, # Counted by uocr blocks_flush below. - "uocr_parked", + ListKind.UOCR_PARKED, # Finished, and held only on the output rank; all ranks must agree. - "deferred", + ListKind.DEFERRED, } if ignore_waiting: - busy_lists.discard("waiting") + busy_lists.discard(ListKind.WAITING) # A PD-decode request waiting with its transferred KV still owns # its request row; a flush would free it. idle &= not any(r.kv.holds_kv for r in self.waiting_queue) diff --git a/python/sglang/srt/managers/scheduler_components/deferred_responses.py b/python/sglang/srt/managers/scheduler_components/deferred_responses.py index 18a208eea32f..21671819cde5 100644 --- a/python/sglang/srt/managers/scheduler_components/deferred_responses.py +++ b/python/sglang/srt/managers/scheduler_components/deferred_responses.py @@ -5,6 +5,7 @@ from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.schedule_batch import FINISH_ABORT +from sglang.srt.managers.scheduler_components.request_list import ListKind if TYPE_CHECKING: from sglang.srt.managers.deferred_output import ResponseCompletion @@ -36,8 +37,8 @@ def __init__( def held_reqs(self) -> List[Req]: return [req for req, _ in self.held.values()] - def request_lists(self) -> Dict[str, List[Req]]: - return {"deferred": self.held_reqs()} + def request_lists(self) -> Dict[ListKind, List[Req]]: + return {ListKind.DEFERRED: self.held_reqs()} def defer(self, req: Req, completion: ResponseCompletion) -> None: if self.disaggregation_mode != DisaggregationMode.NULL: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index cfca970d1884..0b16c0bdfceb 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -26,6 +26,7 @@ kv_private_swa_tokens, kv_private_tokens, ) +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring @@ -70,8 +71,10 @@ class SchedulerInvariantChecker: pool_stats_observer: SchedulerPoolStatsObserver get_last_batch: Callable scheduler_stage_metrics: SchedulerStageMetricsRecorder - # The lists that hold the live requests, by name (`Scheduler.request_lists`). - get_request_lists: Callable[[], Dict[str, List[Req]]] = field(default=lambda: {}) + # The lists that hold the live requests, by kind (`Scheduler.request_lists`). + get_request_lists: Callable[[], Dict[ListKind, Iterable[Req]]] = field( + default=lambda: {} + ) get_chunked_req: Callable = field(default=lambda: None) get_running_reqs: Callable[[], Iterable[Req]] = field(default=lambda: ()) get_prefill_launched_reqs: Callable[[], Iterable[Req]] = field(default=lambda: ()) @@ -484,14 +487,14 @@ def maybe_check_request_lists(self, batch: ScheduleBatch) -> None: live request is in exactly one list, one that matches its state.""" if not envs.SGLANG_ENABLE_RUNNING_REQS_CHECK.get(): return - owners: Dict[Req, str] = {} - for name, reqs in self.get_request_lists().items(): + owners: Dict[Req, ListKind] = {} + for kind, reqs in self.get_request_lists().items(): for req in reqs: if req in owners: raise RuntimeError( - f"Request {req.rid} is in both {owners[req]} and {name}" + f"Request {req.rid} is in both {owners[req].name} and {kind.name}" ) - owners[req] = name + owners[req] = kind chunked_req = self.get_chunked_req() diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index 195a55c15a7d..d0e4db10272a 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -1,16 +1,86 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Dict, Iterable, Iterator, Set +import enum +from abc import ABC, abstractmethod +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + Iterable, + Iterator, + List, + Optional, + Set, +) if TYPE_CHECKING: + from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.schedule_batch import Req -class RequestList: +class ListKind(enum.IntEnum): + """The lists that hold the live requests. Values below 16 are the core's own + lists; the lists of other owners take values from 16 up.""" + + WAITING = 1 + # The running requests. + DECODING = 2 + CHUNKED = 3 + PD_BOOTSTRAP = 16 + # Prefill requests out of the chunked slot whose launched steps are not done. + PD_PREFILLING = 17 + PD_SENDING = 18 + PD_PREALLOC = 24 + PD_RETRACTED = 25 + PD_HELD = 26 + PD_RECEIVING = 27 + PD_RELEASING = 29 + HISPARSE_STAGING = 32 + GRAMMAR = 36 + DLLM = 40 + UOCR_PARKED = 48 + DEFERRED = 56 + + +class RequestList(ABC): + """A list of live requests. Every live request is in exactly one list.""" + + kind: ListKind + + @abstractmethod + def __iter__(self) -> Iterator[Req]: ... + + @abstractmethod + def join(self, items) -> None: ... + + @abstractmethod + def leave(self, reqs: Iterable[Req]) -> None: ... + + def abort(self, recv_req: AbortReq) -> None: + """Abort the requests `recv_req` names.""" + + def take(self, pred: Callable[[Req], bool]) -> List[Req]: + """Drop the requests matching `pred` and return them, in list order.""" + taken = [req for req in self if pred(req)] + self.leave(taken) + return taken + + def __len__(self) -> int: + return sum(1 for _ in self) + + +class _RequestList(RequestList): """Requests in join order, held by identity. Each change rebinds a new dict, so leaving while iterating is safe.""" - def __init__(self) -> None: + def __init__( + self, + kind: ListKind, + *, + abort: Optional[Callable[[AbortReq], None]] = None, + ) -> None: + self.kind = kind + self._abort = abort self._reqs: Dict[Req, None] = {} def join(self, reqs: Iterable[Req]) -> None: @@ -26,6 +96,10 @@ def leave(self, reqs: Iterable[Req]) -> None: def leave_finished(self) -> None: self._drop({req for req in self._reqs if req.finished()}) + def abort(self, recv_req: AbortReq) -> None: + if self._abort is not None: + self._abort(recv_req) + def __iter__(self) -> Iterator[Req]: return iter(self._reqs) @@ -38,3 +112,13 @@ def __contains__(self, req: Req) -> bool: def _drop(self, leaving: Set[Req]) -> None: if leaving: self._reqs = {req: None for req in self._reqs if req not in leaving} + + +def move(reqs: Iterable[Req], src: RequestList, dst: RequestList) -> None: + """Move requests from `src` to `dst`; each must be in `src`.""" + reqs = list(reqs) + missing = [req.rid for req in reqs if req not in src] + if missing: + raise RuntimeError(f"Requests {missing} are not in {src.kind.name}") + src.leave(reqs) + dst.join(reqs) diff --git a/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py b/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py index 500dc17c0e06..71904561c66e 100644 --- a/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py +++ b/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Optional from sglang.srt.managers.overlap_utils import resolve_forward_inputs +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.models.unlimited_ocr_v2_support.async_shm import ( create_async_shm_materializer, ) @@ -102,8 +103,8 @@ def take_released(self) -> List[Req]: """Parked windows whose predecessor finished, to queue in this order.""" return self._prefetch.take_released() - def request_lists(self) -> Dict[str, List[Req]]: - return {"uocr_parked": self._prefetch.held_reqs()} + def request_lists(self) -> Dict[ListKind, List[Req]]: + return {ListKind.UOCR_PARKED: self._prefetch.held_reqs()} def on_req_dropped(self, req: Req) -> None: """Release reuse and prefetch state of a request aborted before running.""" diff --git a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py index 2012a733b4d2..0b25a71a8f90 100644 --- a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py +++ b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py @@ -24,7 +24,10 @@ ) from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams @@ -476,7 +479,7 @@ def test_retry_release_retains_donated_checkpoint(self): processed_tokens_counter=0, _release_aborted_request=lambda req: None, clear_pending_chunk_send=lambda req: None, - prefill_launched_reqs=RequestList(), + prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), ) with patch( "sglang.srt.disaggregation.prefill.get_disagg", diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index ac035e24c803..d149a8f6be60 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -9,7 +9,10 @@ from sglang.srt.disaggregation.utils import MetadataBuffers from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ReqKvInfo -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle from sglang.srt.runtime_context import get_context @@ -57,7 +60,7 @@ def update_finish_state(self): class _Scheduler(SchedulerDisaggregationPrefillMixin): def __init__(self): - self.prefill_launched_reqs = RequestList() + self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) self.batch_result_processor = SimpleNamespace( snapshot_auxiliary_output_starts=Mock(return_value=[]), move_logprobs_to_cpu=Mock(), diff --git a/test/registered/unit/disaggregation/test_prefill_complete_integration.py b/test/registered/unit/disaggregation/test_prefill_complete_integration.py index 065ec46c90a8..20068515d768 100644 --- a/test/registered/unit/disaggregation/test_prefill_complete_integration.py +++ b/test/registered/unit/disaggregation/test_prefill_complete_integration.py @@ -24,7 +24,10 @@ from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.base_prefix_cache import ( CacheRequestHandle, @@ -463,7 +466,7 @@ def test_chunked_prefill_does_not_yield_for_decode_admission(self): check_bootstrap=Mock(return_value=False), has_bootstrapped_waiting_req=Mock(return_value=True), optimistic_release_and_requeue=Mock(), - prefill_launched_reqs=RequestList(), + prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), ) scheduler.checkpoint_disagg_prefill = MethodType( SchedulerDisaggregationPrefillMixin.checkpoint_disagg_prefill, diff --git a/test/registered/unit/disaggregation/test_prefill_result_polling.py b/test/registered/unit/disaggregation/test_prefill_result_polling.py index f78847989c1c..a00e133b6997 100644 --- a/test/registered/unit/disaggregation/test_prefill_result_polling.py +++ b/test/registered/unit/disaggregation/test_prefill_result_polling.py @@ -31,7 +31,10 @@ from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.radix_cache import RadixCache from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries @@ -175,8 +178,8 @@ def __init__( self.chunked_req = None self.decode_offload_manager = None self.running_batch = ScheduleBatch(reqs=[]) - self.running_reqs = RequestList() - self.prefill_launched_reqs = RequestList() + self.running_reqs = _RequestList(ListKind.DECODING) + self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) self.invariant_checker = Mock() self.uocr_runtime = None self.waiting_queue = [] @@ -1262,8 +1265,8 @@ def __init__(self): self.forward_sleep_time = None self.disaggregation_mode = DisaggregationMode.PREFILL self.dllm_config = None - self.running_reqs = RequestList() - self.prefill_launched_reqs = RequestList() + self.running_reqs = _RequestList(ListKind.DECODING) + self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) self.enable_staging = False self.enable_overlap = False self.token_to_kv_pool_allocator = SimpleNamespace(page_size=64) diff --git a/test/registered/unit/lora/test_lora_pp.py b/test/registered/unit/lora/test_lora_pp.py index 85cec5ba5a57..c0ec461e9245 100644 --- a/test/registered/unit/lora/test_lora_pp.py +++ b/test/registered/unit/lora/test_lora_pp.py @@ -33,7 +33,10 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import ( SchedulerMetricsReporter, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.runtime_context import get_context, get_parallel register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -125,7 +128,7 @@ def request(lora_id, max_new_tokens): scheduler.__dict__.update( running_mbs=[first, second], mbs=[first, second], - running_reqs=RequestList(), + running_reqs=_RequestList(ListKind.DECODING), waiting_queue=[waiting], lora_drainer=LoRADrainer(max_loras_per_batch=2, max_wait_time_secs=3), enable_lora=True, diff --git a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py index 87928cc91f32..c5cc22014622 100644 --- a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py +++ b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py @@ -24,6 +24,7 @@ SchedulerPoolStatsObserver, kv_private_swa_tokens, ) +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.allocator.page_interleave import PageInterleavePoolAllocator from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams @@ -205,8 +206,10 @@ def _make_checker(self): pool_stats_observer=observer, get_last_batch=lambda: self.last_batch, get_request_lists=lambda: { - "chunked": [] if self.chunked_req is None else [self.chunked_req], - "running": self.running_batch.reqs, + ListKind.CHUNKED: ( + [] if self.chunked_req is None else [self.chunked_req] + ), + ListKind.DECODING: self.running_batch.reqs, }, scheduler_stage_metrics=None, ) diff --git a/test/registered/unit/managers/test_deferred_output.py b/test/registered/unit/managers/test_deferred_output.py index 20d0634c21c3..4c666541dc28 100644 --- a/test/registered/unit/managers/test_deferred_output.py +++ b/test/registered/unit/managers/test_deferred_output.py @@ -18,7 +18,10 @@ from sglang.srt.managers.scheduler_components.output_streamer import ( SchedulerOutputStreamer, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) @@ -54,7 +57,7 @@ def _scheduler(): scheduler.chunked_req = None scheduler.hisparse_coordinator = None scheduler.collect_inflight_reqs = lambda: set() - scheduler.prefill_launched_reqs = RequestList() + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_disagg_idle_step_counters.py b/test/registered/unit/managers/test_disagg_idle_step_counters.py index b102d6ebe29a..69f852fad413 100644 --- a/test/registered/unit/managers/test_disagg_idle_step_counters.py +++ b/test/registered/unit/managers/test_disagg_idle_step_counters.py @@ -20,7 +20,10 @@ ) from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.runtime_context import get_parallel @@ -543,8 +546,8 @@ def record_result(self, scheduler, batch, result_ts, *, result=None): def make_scheduler(self, schedule): scheduler = Scheduler.__new__(Scheduler) - scheduler.running_reqs = RequestList() - scheduler.prefill_launched_reqs = RequestList() + scheduler.running_reqs = _RequestList(ListKind.DECODING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.tp_size = get_parallel().tp_size scheduler.gracefully_exit = False scheduler._engine_paused = False diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py index 19c1e6beb802..8358fcc4b357 100644 --- a/test/registered/unit/managers/test_kv_page_invariants.py +++ b/test/registered/unit/managers/test_kv_page_invariants.py @@ -7,6 +7,7 @@ import torch from sglang.srt.managers.schedule_batch import ReqKvInfo +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -62,7 +63,9 @@ def __init__(self, rpi, committed, allocated, rid=None): class TestKVPageInvariants(CustomTestCase): def test_committed_gt_allocated_raises(self): chk, rtt, tc, alloc = _make_checker() - chk.get_request_lists = lambda: {"running": [_FakeOwner(0, 145, 144, rid="a")]} + chk.get_request_lists = lambda: { + ListKind.DECODING: [_FakeOwner(0, 145, 144, rid="a")] + } with self.assertRaises(AssertionError): chk._check_kv_page_invariants() @@ -79,13 +82,15 @@ def test_owner_references_free_page_raises(self): rtt[0, :3] = torch.tensor( [5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2] ) - chk.get_request_lists = lambda: {"running": [_FakeOwner(0, 3, 3, rid="a")]} + chk.get_request_lists = lambda: { + ListKind.DECODING: [_FakeOwner(0, 3, 3, rid="a")] + } with self.assertRaises(ValueError): chk._check_kv_page_invariants() def test_requests_in_any_list_are_checked(self): # A request still owns its row while running or parked between chunks. - for where in ("running", "chunked"): + for where in (ListKind.DECODING, ListKind.CHUNKED): with self.subTest(where=where): chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([5, 6, 7])) rtt[0, :3] = torch.tensor( @@ -104,7 +109,7 @@ def test_classed_allocator_without_flat_free_list_is_checked(self): ) rtt[0, :1] = torch.tensor([5 * alloc.page_size]) chk.get_request_lists = lambda: { - "running": [_FakeOwner(0, 1, 1, rid="classed")] + ListKind.DECODING: [_FakeOwner(0, 1, 1, rid="classed")] } with patch( @@ -119,7 +124,9 @@ def test_classed_allocator_without_flat_free_list_is_checked(self): def test_free_pool_duplicate_raises(self): chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4])) rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE]) # owner page 10, not in free - chk.get_request_lists = lambda: {"running": [_FakeOwner(0, 1, 1, rid="a")]} + chk.get_request_lists = lambda: { + ListKind.DECODING: [_FakeOwner(0, 1, 1, rid="a")] + } with self.assertRaises(ValueError): chk._check_kv_page_invariants() diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index 4e9fe6dbb9bc..84315df5f486 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -13,7 +13,10 @@ from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.schedule_batch import NextBatchPlan, Req, ReqKvInfo from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -95,7 +98,7 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler: s.tree_cache = tree_cache s.chunked_req = chunked_req s._pending_chunked_abort_req = None - s.running_reqs = RequestList() + s.running_reqs = _RequestList(ListKind.DECODING) s.invariant_checker = MagicMock() return s diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py index 3be74c493b1b..42722c242442 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py +++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py @@ -12,7 +12,10 @@ from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) @@ -60,7 +63,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.enable_hierarchical_cache = False scheduler.enable_lmcache = False scheduler.enable_hisparse = False - scheduler.prefill_launched_reqs = RequestList() + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_scheduler_held_row_admission.py b/test/registered/unit/managers/test_scheduler_held_row_admission.py index 55428332bf08..d9863273f436 100644 --- a/test/registered/unit/managers/test_scheduler_held_row_admission.py +++ b/test/registered/unit/managers/test_scheduler_held_row_admission.py @@ -13,7 +13,10 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.runtime_context import get_context register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -100,7 +103,7 @@ def _scheduler(self, disaggregation_mode, running, final_chunk, waiting): available_size=lambda: NUM_ROWS - len(running) - 1 ) s.beam_coordinator = SimpleNamespace(pending_member_rows=lambda batch: 0) - s.running_reqs = RequestList() + s.running_reqs = _RequestList(ListKind.DECODING) return s def _admitted(self, disaggregation_mode): diff --git a/test/registered/unit/managers/test_scheduler_hicache_events.py b/test/registered/unit/managers/test_scheduler_hicache_events.py index eed3ca937203..7f8cd36aed2d 100644 --- a/test/registered/unit/managers/test_scheduler_hicache_events.py +++ b/test/registered/unit/managers/test_scheduler_hicache_events.py @@ -15,7 +15,10 @@ from sglang.srt.managers.schedule_batch import NextBatchPlan from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -33,7 +36,7 @@ def setUp(self): s._process_storage_prefetch_retries = self.calls.retry s.process_pending_chunked_abort = Mock() s.process_prefill_chunk = Mock() - s.running_reqs = RequestList() + s.running_reqs = _RequestList(ListKind.DECODING) s.invariant_checker = Mock() s.dp_attn_adapter = Mock() s.dp_attn_adapter.maybe_prepare_mlp_sync_batch.return_value = None diff --git a/test/registered/unit/managers/test_scheduler_on_idle_load.py b/test/registered/unit/managers/test_scheduler_on_idle_load.py index 8fe82038efa2..cfb78cb2b057 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/test/registered/unit/managers/test_scheduler_on_idle_load.py @@ -23,7 +23,10 @@ from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) @@ -120,7 +123,7 @@ def _idle_scheduler(self, materializer): scheduler.grammar_manager = SimpleNamespace(grammar_queue=[]) scheduler.disaggregation_mode = DisaggregationMode.NULL scheduler.enable_hisparse = False - scheduler.prefill_launched_reqs = RequestList() + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 2d92a7d558d7..0e4d3a3bcfae 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -33,7 +33,10 @@ SchedulerKvEventsPublisher, ) from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) @@ -98,8 +101,8 @@ def _new_scheduler(self) -> Scheduler: scheduler.last_batch = None scheduler.cur_batch_for_debug = None scheduler.chunked_req = None - scheduler.running_reqs = RequestList() - scheduler.prefill_launched_reqs = RequestList() + scheduler.running_reqs = _RequestList(ListKind.DECODING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.running_batch = MagicMock() scheduler.running_batch.reqs = [] scheduler.running_batch.is_empty.return_value = True diff --git a/test/registered/unit/managers/test_scheduler_running_reqs.py b/test/registered/unit/managers/test_scheduler_running_reqs.py index 965a0ef13224..71381297f17c 100644 --- a/test/registered/unit/managers/test_scheduler_running_reqs.py +++ b/test/registered/unit/managers/test_scheduler_running_reqs.py @@ -31,7 +31,10 @@ from sglang.srt.managers.scheduler_components.invariant_checker import ( SchedulerInvariantChecker, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -95,7 +98,7 @@ def _empty_request_lists(s): s.waiting_queue = [] s.grammar_manager = SimpleNamespace(grammar_queue=[]) s.chunked_req = None - s.prefill_launched_reqs = RequestList() + s.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) s.disagg_prefill_transfer_queue = None s.disagg_prefill_bootstrap_queue = None s.disagg_decode_prealloc_queue = None @@ -113,7 +116,7 @@ def _empty_request_lists(s): def _scheduler(): s = _Scheduler.__new__(_Scheduler) s.running_batch = ScheduleBatch(reqs=[]) - s.running_reqs = RequestList() + s.running_reqs = _RequestList(ListKind.DECODING) s.invariant_checker = _invariant_checker(s) s.disaggregation_mode = DisaggregationMode.NULL s.dllm_config = None @@ -167,7 +170,7 @@ def _launch(s, batch): class TestRequestList(CustomTestCase): def test_join_keeps_order_and_leaving_mid_iteration_is_safe(self): a, b, c, outsider = (_req(rid) for rid in "abcx") - running = RequestList() + running = _RequestList(ListKind.DECODING) running.join([a, b]) running.join([c, b]) _finish(b) @@ -200,7 +203,7 @@ def test_fails_a_request_in_two_lists(self): self.s.running_reqs.join([req]) self.s.waiting_queue.append(req) - with self.assertRaisesRegex(RuntimeError, "a is in both waiting and running"): + with self.assertRaisesRegex(RuntimeError, "a is in both WAITING and DECODING"): self.check(ScheduleBatch(reqs=[])) def test_fails_a_launched_row_in_no_list(self): @@ -319,7 +322,7 @@ def test_ngram_grammar_step_launches_no_finished_row(self): s = Scheduler.__new__(Scheduler) s.running_batch = running s.last_batch = None - s.running_reqs = RequestList() + s.running_reqs = _RequestList(ListKind.DECODING) s.running_reqs.join(running.reqs) s.disaggregation_mode = DisaggregationMode.NULL s.dllm_config = None diff --git a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py index c633d75645ce..3db4629fe735 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -17,7 +17,10 @@ from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) -from sglang.srt.managers.scheduler_components.request_list import RequestList +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_components.weight_updater import ( SchedulerWeightUpdaterManager, _WeightUpdateSession, @@ -169,8 +172,8 @@ def _scheduler(self) -> Scheduler: scheduler.load_publisher = MagicMock() scheduler.load_inquirer = MagicMock() scheduler.tree_cache = SimpleNamespace(flush_pending_backups=MagicMock()) - scheduler.running_reqs = RequestList() - scheduler.prefill_launched_reqs = RequestList() + scheduler.running_reqs = _RequestList(ListKind.DECODING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) scheduler.invariant_checker = MagicMock() scheduler.kv_weight_version_tracker = KvWeightVersionTracker( num_slots=8, From 72626bb5c6fc398b6d44f9f177a512b7fe065c12 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:11:07 -0800 Subject: [PATCH 03/12] migrate grammar, pd prefill bootstrap and transfer, pd decode retracted and held, and dllm to request lists; abort through each list --- .../sglang/srt/constrained/grammar_manager.py | 33 ++++--- python/sglang/srt/disaggregation/decode.py | 70 +++++++++------ python/sglang/srt/disaggregation/prefill.py | 90 +++++++++++++------ .../sglang/srt/disaggregation/role_switch.py | 3 +- python/sglang/srt/dllm/mixin/scheduler.py | 70 +++++++++------ python/sglang/srt/managers/schedule_policy.py | 8 +- python/sglang/srt/managers/scheduler.py | 67 ++------------ python/sglang/test/scheduler_request_lists.py | 50 +++++++++++ .../test_disaggregation_optimistic_prefill.py | 2 +- .../unit/constrained/test_grammar_manager.py | 34 +++---- .../test_admission_abort_not_enqueued.py | 18 ++-- .../test_prefill_abort_result_cleanup.py | 9 +- .../test_prefill_complete_integration.py | 62 ++++++++----- .../test_prefill_result_polling.py | 25 ++++-- .../unit/managers/test_auxiliary_output.py | 12 ++- .../unit/managers/test_deferred_output.py | 2 +- .../managers/test_multimodal_abort_cleanup.py | 13 +-- ...test_priority_scheduling_disaggregation.py | 28 +++--- .../test_scheduler_chunked_abort_race.py | 14 +-- ...est_scheduler_flush_cache_after_retract.py | 2 +- .../managers/test_scheduler_on_idle_load.py | 2 +- .../test_scheduler_pause_generation.py | 18 ++-- .../managers/test_scheduler_running_reqs.py | 32 +------ .../unit/managers/test_scheduler_timeouts.py | 10 +-- .../test_decode_retraction_backup.py | 9 +- .../unit/mem_cache/test_dllm_fdfo_kv_reuse.py | 28 +++--- .../unit/mem_cache/test_mamba_capacity.py | 12 ++- 27 files changed, 402 insertions(+), 321 deletions(-) create mode 100644 python/sglang/test/scheduler_request_lists.py diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 924e0bcfb013..237e7f7db45d 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/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_requests) if not get_serving().skip_tokenizer_init: self.grammar_backend = create_grammar_backend( self.server_args, @@ -112,6 +116,9 @@ def _pp_sync_ready_failed( return data def abort_requests(self, recv_req: AbortReq): + # Abort method 2: call `set_finish_with_abort` + # The request will still run one prefill forward pass. + # In this case, we change the input_ids to be only one token to make this prefill cheap. for req in self.grammar_queue: if recv_req.abort_all or req.rid.startswith(recv_req.rid): logger.debug(f"Abort grammar queue request. {req.rid=}") @@ -204,7 +211,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,6 +231,8 @@ 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() @@ -231,7 +240,7 @@ def get_ready_grammar_requests(self) -> List[Req]: # 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 +254,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 +305,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 +327,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 +339,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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 4e10a64ff80f..401de66a29a3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -90,8 +90,15 @@ kv_private_swa_tokens, kv_private_tokens, ) -from sglang.srt.managers.scheduler_components.request_list import ListKind -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 @@ -489,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[ @@ -502,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( @@ -919,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 @@ -1010,12 +1021,12 @@ 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 = list(self.held_rebootstrap_reqs) + self.held_rebootstrap_reqs.leave(held) for req in held: self.add(req, is_rebootstrap=True) @@ -1056,17 +1067,28 @@ 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.abort_all or req.rid.startswith(recv_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.abort_all or req.rid.startswith(recv_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[ListKind, List[Req]]: return { @@ -1090,7 +1112,7 @@ def release_memory_occupation(self): self.tree_cache, get_disagg().disaggregation_decode_retraction_backup, ) - self.retracted_queue.clear() + self.retracted_queue.leave(list(self.retracted_queue)) if hasattr(self.kv_manager, "deregister_buffer_to_engine"): self.kv_manager.deregister_buffer_to_engine() @@ -1105,7 +1127,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 = ( @@ -1116,7 +1137,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 @@ -1149,7 +1170,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( @@ -1176,11 +1196,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 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 705ed2cc02bf..4722a7caf29d 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -77,6 +77,12 @@ Req, ScheduleBatch, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + RequestList, + _RequestList, + move, +) from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome from sglang.srt.mem_cache.common import ( checkpoint_kv_cache, @@ -106,6 +112,7 @@ from sglang.srt.utils.common import async_d2h if TYPE_CHECKING: + from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.memory_pool import KVCache @@ -206,7 +213,7 @@ def __init__( self.pp_size = parallel.pp_size self.gpu_id = gpu_id self.bootstrap_port = bootstrap_port - self.queue: List[Req] = [] + self.queue = _RequestList(ListKind.PD_BOOTSTRAP, abort=self.abort_requests) self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics # Each rank owns a subset of logical KV pages, so capacity widens by N. @@ -468,12 +475,24 @@ def add(self, req: Req, num_kv_heads: int) -> None: return if not self.create_sender(req, num_kv_heads): return - self.queue.append(req) + self.queue.join([req]) def extend(self, reqs: List[Req], num_kv_heads: int) -> None: for req in reqs: self.add(req, num_kv_heads) + def abort_requests(self, recv_req: AbortReq) -> None: + # Abort requests that have not yet been bootstrapped + for req in self.queue: + if recv_req.abort_all or req.rid.startswith(recv_req.rid): + logger.debug(f"Abort bootstrap queue request. {req.rid=}") + self.scheduler._release_aborted_request(req) + + if hasattr(req.disagg_kv_sender, "abort"): + req.disagg_kv_sender.abort() + if get_parallel().pp_size > 1: + prepare_abort(req, "Aborted by AbortReq.") + def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: if len(req.origin_input_ids) > self.max_input_len: message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_input_len}" @@ -512,9 +531,11 @@ def pop_bootstrapped(self) -> List[Req]: if len(self.queue) == 0: return [] + # The polls are reduced across ranks in this order. + queue = list(self.queue) if self.pp_size > 1: polls = poll_and_all_reduce_prefill_pp( - [req.disagg_kv_sender for req in self.queue], + [req.disagg_kv_sender for req in queue], self.scheduler.attn_cp_cpu_group, self.scheduler.attn_tp_cpu_group, self.scheduler.pp_group, @@ -524,12 +545,12 @@ def pop_bootstrapped(self) -> List[Req]: ) else: polls = poll_and_all_reduce_attn_cp_tp_group( - [req.disagg_kv_sender for req in self.queue], + [req.disagg_kv_sender for req in queue], self.scheduler.attn_cp_cpu_group, self.scheduler.attn_tp_cpu_group, ) - for i, (req, poll) in enumerate(zip(self.queue, polls)): + for i, (req, poll) in enumerate(zip(queue, polls)): if poll == KVPoll.Failed: self.scheduler.handle_bootstrap_failure(req) indices_to_remove.add(i) @@ -560,14 +581,12 @@ def pop_bootstrapped(self) -> List[Req]: f"Unexpected poll state {poll} for req {req.rid} in pop_bootstrapped" ) - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue.leave([queue[i] for i in indices_to_remove]) return bootstrapped_reqs def release_memory_occupation(self): - self.queue.clear() + self.queue.leave(list(self.queue)) if hasattr(self.kv_manager, "deregister_buffer_to_engine"): self.kv_manager.deregister_buffer_to_engine() @@ -1036,7 +1055,7 @@ def advance_logprob_pt(i: int, req: Req) -> None: # Test hook: exercise the release/requeue retry path. if req.pending_bootstrap and should_force_retry(req): - self.optimistic_release_and_requeue(req) + self.optimistic_release_and_requeue(req, self.prefill_launched_reqs) advance_logprob_pt(i, req) continue @@ -1084,8 +1103,11 @@ def advance_logprob_pt(i: int, req: Req) -> None: i, req, logits_output ) self.checkpoint_disagg_prefill(req) - self.prefill_launched_reqs.leave([req]) - self.disagg_prefill_transfer_queue.append(req) + move( + [req], + self.prefill_launched_reqs, + self.disagg_prefill_transfer_queue, + ) if self.spec_algorithm.is_eagle() and draft_input is not None: req.output_topk_p = draft_input.topk_p[i] req.output_topk_index = draft_input.topk_index[i] @@ -1145,7 +1167,7 @@ def advance_logprob_pt(i: int, req: Req) -> None: # Only a yield leaves the chunked slot before the final chunk; the # yield is not undone by a bootstrap that completes meanwhile. if not still_chunking: - self.optimistic_release_and_requeue(req) + self.optimistic_release_and_requeue(req, self.prefill_launched_reqs) advance_logprob_pt(i, req) req.time_stats.set_last_chunked_prefill_finish_time() continue @@ -1196,6 +1218,14 @@ def advance_logprob_pt(i: int, req: Req) -> None: ) self.maybe_send_health_check_signal() + def abort_disagg_prefill_transfer_reqs(self: Scheduler, recv_req: AbortReq) -> None: + # Abort in-flight requests + for req in self.disagg_prefill_transfer_queue: + if recv_req.abort_all or req.rid.startswith(recv_req.rid): + logger.debug(f"Abort transfer queue request. {req.rid=}") + if hasattr(req.disagg_kv_sender, "abort"): + req.disagg_kv_sender.abort() + @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE) def process_disagg_prefill_transfer_queue( self: Scheduler, rids_to_check: Optional[Dict[str, bool]] = None @@ -1211,15 +1241,17 @@ def process_disagg_prefill_transfer_queue( done_reqs = [] + # The polls are reduced across ranks in this order. + transfer_queue = list(self.disagg_prefill_transfer_queue) polls = poll_and_all_reduce_attn_cp_tp_group( - [req.disagg_kv_sender for req in self.disagg_prefill_transfer_queue], + [req.disagg_kv_sender for req in transfer_queue], self.attn_cp_cpu_group, self.attn_tp_cpu_group, ) undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_transfer_queue. If Success, respond to the client and remove it from the queue - for req, poll in zip(self.disagg_prefill_transfer_queue, polls): + for req, poll in zip(transfer_queue, polls): if rids_to_check is not None: if req.rid not in rids_to_check: undone_reqs.append(req) @@ -1271,7 +1303,9 @@ def process_disagg_prefill_transfer_queue( # Decode may have admitted different requests. Yield one # slot so they can prefill, including peers still waiting # for a metadata buffer in the bootstrap queue. - self.optimistic_release_and_requeue(req) + self.optimistic_release_and_requeue( + req, self.disagg_prefill_transfer_queue + ) elif poll != KVPoll.Failed: undone_reqs.append(req) continue @@ -1335,7 +1369,7 @@ def process_disagg_prefill_transfer_queue( req, self.req_to_metadata_buffer_idx_allocator ) - self.disagg_prefill_transfer_queue = undone_reqs + self.disagg_prefill_transfer_queue.leave(set(transfer_queue) - set(undone_reqs)) return done_reqs @@ -1490,13 +1524,14 @@ def process_prefill_chunk( != "prefill_complete" and self.has_bootstrapped_waiting_req() ): - # optimistic request yields to waiting requests + # optimistic request yields to waiting requests; with a step + # in flight, it is requeued once its last step is processed. self.chunked_req = None + self.prefill_launched_reqs.join([req]) if req.inflight_steps == 0: - self.optimistic_release_and_requeue(req) - else: - # Requeued once its last step in flight is processed. - self.prefill_launched_reqs.join([req]) + self.optimistic_release_and_requeue( + req, self.prefill_launched_reqs + ) # else: still bootstrapping, keep computing without sending elif self.enable_overlap: # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved @@ -1883,8 +1918,10 @@ def _request_row_payload(): else: self.disagg_prefill_pending_chunk_rids.add(req.rid) - def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None: - """Release KV cache and requeue an optimistic prefill request.""" + def optimistic_release_and_requeue( + self: Scheduler, req: Req, src: RequestList + ) -> None: + """Release KV cache and requeue an optimistic prefill request out of `src`.""" max_attempts = get_disagg().optimistic_prefill_attempts uses_write_through_cache = _uses_write_through_cache(self.tree_cache) if not uses_write_through_cache: @@ -1936,8 +1973,7 @@ def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None: maybe_release_metadata_buffer( req, self.req_to_metadata_buffer_idx_allocator ) - self.prefill_launched_reqs.leave([req]) - self.disagg_prefill_bootstrap_queue.queue.append(req) + move([req], src, self.disagg_prefill_bootstrap_queue.queue) else: req.prefill_attempt_count += 1 logger.info( @@ -1948,5 +1984,5 @@ def optimistic_release_and_requeue(self: Scheduler, req: Req) -> None: self.metrics_collector.increment_prefill_retries(1) req.time_stats.set_wait_queue_entry_time() req.arrival_processed_tokens = self.processed_tokens_counter - self.prefill_launched_reqs.leave([req]) + src.leave([req]) self.waiting_queue.insert(0, req) diff --git a/python/sglang/srt/disaggregation/role_switch.py b/python/sglang/srt/disaggregation/role_switch.py index a7d91f276c2c..781873bf9b3b 100644 --- a/python/sglang/srt/disaggregation/role_switch.py +++ b/python/sglang/srt/disaggregation/role_switch.py @@ -247,7 +247,8 @@ def teardown_disaggregation(scheduler: Scheduler) -> None: if km is not None: km.teardown() scheduler.disagg_prefill_bootstrap_queue = None - scheduler.disagg_prefill_transfer_queue = [] + transfer_queue = scheduler.disagg_prefill_transfer_queue + transfer_queue.leave(list(transfer_queue)) elif mode == DisaggregationMode.DECODE: q = getattr(scheduler, "disagg_decode_prealloc_queue", None) if q is not None: diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 064fb051dd89..8647c5576e4d 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -2,12 +2,22 @@ import logging from array import array -from typing import TYPE_CHECKING, List, Optional, Set, Union +from typing import TYPE_CHECKING, Callable, List, Optional, Set, Union from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.mixin.req import DllmReqPhase -from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import ( + FINISH_ABORT, + FINISH_LENGTH, + Req, + ScheduleBatch, +) from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) +from sglang.srt.managers.utils import make_abort_req from sglang.srt.mem_cache.common import release_kv_cache from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.observability.req_time_stats import set_time_batch @@ -16,6 +26,7 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: + from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.utils import GenerationBatchResult @@ -27,7 +38,20 @@ def init_diffusion_llm(self: Scheduler): if get_exec().dllm.dllm_algorithm is not None else None ) - self.dllm_manager = DllmManager(dllm_config=self.dllm_config) + self.dllm_manager = DllmManager( + dllm_config=self.dllm_config, + release_aborted=self._release_aborted_dllm_req, + ) + + def _release_aborted_dllm_req(self: Scheduler, req: Req) -> None: + # Finished in place, so a later abort in the same intake skips it. + if req.finished_reason is None: + req.finish_outside_step(FINISH_ABORT()) + self._release_aborted_request(req) + self.send_abort_output(make_abort_req(req), req) + if req.kv.holds_kv or req.kv.holds_mamba: + release_kv_cache(req, self.tree_cache, checkpoint=False) + logger.debug(f"Abort dLLM queued request. {req.rid=}") def validate_dllm_request(self: Scheduler, req: Req) -> Optional[str]: if self.dllm_config is None: @@ -424,13 +448,19 @@ class DllmManager: - staging_queue: Requests allocated resources by PrefillAdder """ - def __init__(self, dllm_config: Optional[DllmConfig] = None): + def __init__( + self, + dllm_config: Optional[DllmConfig], + release_aborted: Callable[[Req], None], + ): self.dllm_config = dllm_config self.max_running_reqs = ( dllm_config.max_running_requests if dllm_config is not None else 1 ) - self.waiting_queue: List[Req] = [] + self.waiting_queue = _RequestList(ListKind.DLLM, abort=self._abort) + # A subset of the waiting queue. self.staging_queue: List[Req] = [] + self._release_aborted = release_aborted def get_prefill_requests(self) -> List[Req]: """Get all prefill requests from waiting queue.""" @@ -450,7 +480,7 @@ def add_waiting_reqs(self, reqs: Union[Req, List[Req]]) -> None: if self._has_duplicate_reqs(reqs_to_add): raise RuntimeError("Redundant requests detected in dLLM requests.") - self.waiting_queue.extend(reqs_to_add) + self.waiting_queue.join(reqs_to_add) def add_staging_reqs(self, reqs: Union[Req, List[Req]]) -> None: """Add requests to staging queue (allocated by PrefillAdder).""" @@ -474,27 +504,17 @@ def is_empty(self) -> bool: def filter_finished_reqs(self) -> None: """Remove finished requests from both queues.""" - self.waiting_queue = [req for req in self.waiting_queue if not req.finished()] + self.waiting_queue.leave_finished() self.staging_queue = [req for req in self.staging_queue if not req.finished()] - def pop_aborted_reqs(self, abort_all: bool, rid: str) -> List[Req]: - aborted_reqs: List[Req] = [] - seen: Set[int] = set() - - for queue_name in ("waiting_queue", "staging_queue"): - queue = getattr(self, queue_name) - kept_queue = [] - for req in queue: - if abort_all or req.rid.startswith(rid): - req_id = id(req) - if req_id not in seen: - aborted_reqs.append(req) - seen.add(req_id) - else: - kept_queue.append(req) - setattr(self, queue_name, kept_queue) - - return aborted_reqs + def _abort(self, recv_req: AbortReq) -> None: + aborted_reqs = self.waiting_queue.take( + lambda req: recv_req.abort_all or req.rid.startswith(recv_req.rid) + ) + aborted = set(aborted_reqs) + self.staging_queue = [req for req in self.staging_queue if req not in aborted] + for req in aborted_reqs: + self._release_aborted(req) def init_next_round(self) -> None: """Initialize staging requests for next round and clear staging queue.""" diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 89661d243cd5..6aa07e064ff1 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -44,7 +44,7 @@ from dataclasses import dataclass from enum import Enum, auto from functools import lru_cache -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Set, Union import torch @@ -207,7 +207,7 @@ def calc_priority( waiting_queue: List[Req], running_batch: Optional[ScheduleBatch] = None, processed_tokens: int = 0, - bootstrap_queue: List[Req] = [], + bootstrap_queue: Iterable[Req] = (), ) -> None: policy = self._determine_active_policy(waiting_queue) @@ -276,7 +276,7 @@ def _touch_waiting_prefixes( self, policy: Policy, waiting_queue: List[Req], - bootstrap_queue: List[Req] = [], + bootstrap_queue: Iterable[Req] = (), ) -> None: # Head last, so LRU evicts in reverse admission order; cache-aware policies # already touch while matching, and other eviction strategies ignore recency. @@ -298,7 +298,7 @@ def _touch_waiting_prefixes( if not refresh: return # Bootstrap-queued requests run after the waiting queue, so touch them first. - for r in reversed(bootstrap_queue[:bootstrap_max]): + for r in reversed(list(bootstrap_queue)[:bootstrap_max]): touch_waiting_prefix(r, self.tree_cache) for r in reversed(waiting_queue[:WAITING_QUEUE_PREFIX_MATCH_MAX]): touch_waiting_prefix(r, self.tree_cache) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index eed16311f4b3..8617d1535357 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -301,6 +301,7 @@ ) from sglang.srt.managers.scheduler_components.request_list import ( ListKind, + RequestList, _RequestList, ) from sglang.srt.managers.scheduler_components.request_receiver import ( @@ -1573,7 +1574,9 @@ def init_watch_dog_memory_saver_input_blocker(self): def init_disaggregation(self): self.mm_receiver = None self.disagg_prefill_bootstrap_queue = None - self.disagg_prefill_transfer_queue = None + self.disagg_prefill_transfer_queue = _RequestList( + ListKind.PD_SENDING, abort=self.abort_disagg_prefill_transfer_reqs + ) self.disagg_decode_prealloc_queue = None self.disagg_decode_transfer_queue = None @@ -1702,7 +1705,6 @@ def init_disaggregation(self): scheduler_stage_metrics=self.scheduler_stage_metrics, transfer_backend=self.transfer_backend, ) - self.disagg_prefill_transfer_queue: List[Req] = [] # Requests with a sent chunk that are not yet on the transfer queue. self.disagg_prefill_pending_chunk_rids: Set[str] = set() @@ -4693,8 +4695,7 @@ def request_lists(self) -> Dict[ListKind, Iterable[Req]]: # Read at call time: a PD role switch rebuilds the PD queues. if self.disagg_prefill_bootstrap_queue is not None: lists[ListKind.PD_BOOTSTRAP] = self.disagg_prefill_bootstrap_queue.queue - if self.disagg_prefill_transfer_queue is not None: - lists[ListKind.PD_SENDING] = self.disagg_prefill_transfer_queue + lists[ListKind.PD_SENDING] = self.disagg_prefill_transfer_queue if self.dllm_config is not None: # Its staging queue is the part of this queue the current round runs. lists[ListKind.DLLM] = self.dllm_manager.waiting_queue @@ -6077,46 +6078,12 @@ def abort_request(self, recv_req: AbortReq): release_kv_cache(req, self.tree_cache, checkpoint=False) logger.debug(f"Abort queued request. {req.rid=}") - if self.dllm_config is not None: - for req in self.dllm_manager.pop_aborted_reqs( - recv_req.abort_all, recv_req.rid - ): - # Finished in place, so a later abort in the same intake skips it. - if req.finished_reason is None: - req.finish_outside_step(FINISH_ABORT()) - self._release_aborted_request(req) - self.send_abort_output(make_abort_req(req), req) - if req.kv.holds_kv or req.kv.holds_mamba: - release_kv_cache(req, self.tree_cache, checkpoint=False) - logger.debug(f"Abort dLLM queued request. {req.rid=}") - - # Delete the requests in the grammar queue - # Abort method 2: call `set_finish_with_abort` - # The request will still run one prefill forward pass. - # In this case, we change the input_ids to be only one token to make this prefill cheap. - self.grammar_manager.abort_requests(recv_req) + for request_list in self.request_lists().values(): + if isinstance(request_list, RequestList): + request_list.abort(recv_req) # Delete requests not in the waiting queue when PD disaggregation is enabled - if self.disaggregation_mode == DisaggregationMode.PREFILL: - # Abort requests that have not yet been bootstrapped - for req in self.disagg_prefill_bootstrap_queue.queue: - if recv_req.abort_all or req.rid.startswith(recv_req.rid): - logger.debug(f"Abort bootstrap queue request. {req.rid=}") - self._release_aborted_request(req) - - if hasattr(req.disagg_kv_sender, "abort"): - req.disagg_kv_sender.abort() - if get_parallel().pp_size > 1: - prepare_abort(req, "Aborted by AbortReq.") - - # Abort in-flight requests - for req in self.disagg_prefill_transfer_queue: - if recv_req.abort_all or req.rid.startswith(recv_req.rid): - logger.debug(f"Abort transfer queue request. {req.rid=}") - if hasattr(req.disagg_kv_sender, "abort"): - req.disagg_kv_sender.abort() - - elif self.disaggregation_mode == DisaggregationMode.DECODE: + if self.disaggregation_mode == DisaggregationMode.DECODE: # Abort requests that have not yet finished preallocation for decode_req in self.disagg_decode_prealloc_queue.queue: if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): @@ -6138,22 +6105,6 @@ def abort_request(self, recv_req: AbortReq): # an ack racing back is never dropped. decode_req.kv_receiver.abort() - # Retracted requests hold a KV backup; held ones hold no KV. - retracted, held = self.disagg_decode_prealloc_queue.take_aborted(recv_req) - for req in retracted: - discard_kv_cache_backup( - req, - self.tree_cache, - get_disagg().disaggregation_decode_retraction_backup, - ) - self.send_abort_output(make_abort_req(req), req) - for req in held: - req.finish_outside_step(FINISH_ABORT()) - self._release_aborted_request(req) - self.send_abort_output( - make_abort_req(req, finished_reason=recv_req.finished_reason), req - ) - # Delete requests in the running batch for req in self.collect_inflight_reqs(): if not req.finished() and ( diff --git a/python/sglang/test/scheduler_request_lists.py b/python/sglang/test/scheduler_request_lists.py new file mode 100644 index 000000000000..2d777c845dbb --- /dev/null +++ b/python/sglang/test/scheduler_request_lists.py @@ -0,0 +1,50 @@ +"""The request lists of a partial scheduler, for tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +from sglang.srt.disaggregation.mode import DisaggregationMode +from sglang.srt.managers.scheduler_components.deferred_responses import ( + DeferredResponses, +) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) + +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler + + +class FakeReq(SimpleNamespace): + """A stand-in request a request list can hold: compared and hashed by + identity, as `Req` is.""" + + __eq__ = object.__eq__ + __hash__ = object.__hash__ + + +def empty_request_lists(scheduler: Scheduler) -> None: + """Give a scheduler built with `Scheduler.__new__` every request list, empty; + the test then fills the ones it needs.""" + scheduler.waiting_queue = [] + scheduler.grammar_manager = SimpleNamespace( + grammar_queue=_RequestList(ListKind.GRAMMAR) + ) + scheduler.chunked_req = None + scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.running_reqs = _RequestList(ListKind.DECODING) + scheduler.disagg_prefill_bootstrap_queue = None + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) + scheduler.dllm_config = None + scheduler.disagg_decode_prealloc_queue = None + scheduler.disagg_decode_transfer_queue = None + scheduler.hisparse_coordinator = None + scheduler.uocr_runtime = None + scheduler.deferred_responses = DeferredResponses( + disaggregation_mode=DisaggregationMode.NULL, + output_streamer=None, + send_abort=None, + ) diff --git a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py index 0b25a71a8f90..64a4b21cd6b9 100644 --- a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py +++ b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py @@ -486,7 +486,7 @@ def test_retry_release_retains_donated_checkpoint(self): return_value=SimpleNamespace(optimistic_prefill_attempts=3), ): SchedulerDisaggregationPrefillMixin.optimistic_release_and_requeue( - scheduler, req + scheduler, req, scheduler.prefill_launched_reqs ) self.assertEqual(tree.total_size(), (self.TRACK_SEQLEN, 1)) diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index 2b2fc73af1fb..4183f76b29bf 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -485,7 +485,7 @@ def test_abort_by_rid_prefix(self): req = _make_req(rid="req-123") future = MagicMock(spec=Future) req.grammar = future - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) abort_req = MagicMock() abort_req.abort_all = False @@ -499,7 +499,7 @@ def test_abort_non_matching_rid(self): mgr = self._make_mgr_with_queue() req = _make_req(rid="req-999") req.grammar = MagicMock(spec=Future) - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) abort_req = MagicMock() abort_req.abort_all = False @@ -514,7 +514,7 @@ def test_abort_all(self): for i in range(3): req = _make_req(rid=f"req-{i}") req.grammar = MagicMock(spec=Future) - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) reqs.append(req) abort_req = MagicMock() @@ -538,7 +538,7 @@ def test_abort_prefix_match(self): mgr = self._make_mgr_with_queue() req = _make_req(rid="req-123-suffix") req.grammar = MagicMock(spec=Future) - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) abort_req = MagicMock() abort_req.abort_all = False @@ -578,7 +578,7 @@ def test_ready_future_returns_req(self): req = _make_req(json_schema="schema") req.grammar = future req.grammar_key = ("json", "schema") - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) result = mgr.get_ready_grammar_requests() self.assertEqual(len(result), 1) @@ -601,7 +601,7 @@ def test_invalid_grammar_aborts_req(self): req = _make_req(json_schema="bad") req.grammar = future req.grammar_key = ("json", "bad") - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) result = mgr.get_ready_grammar_requests() self.assertEqual(len(result), 1) @@ -614,7 +614,7 @@ def test_aborted_req_removed_from_queue(self): req = _make_req(json_schema="schema") req.finished.return_value = True # Already aborted req.grammar = None - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) result = mgr.get_ready_grammar_requests() self.assertEqual(len(result), 1) @@ -629,7 +629,7 @@ def test_timeout_aborts_req(self): req.grammar = future req.grammar_key = ("json", "slow") req.grammar_wait_ct = 0 - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) # First call: not ready, increments wait_ct to 1 (== max_poll) result = mgr.get_ready_grammar_requests() @@ -653,7 +653,7 @@ def test_pending_future_stays_in_queue(self): req.grammar = future req.grammar_key = ("json", "pending") req.grammar_wait_ct = 0 - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) result = mgr.get_ready_grammar_requests() self.assertEqual(len(result), 0) @@ -680,13 +680,13 @@ def test_mixed_ready_and_pending(self): pending_req.grammar_key = ("json", "pending") pending_req.grammar_wait_ct = 0 - mgr.grammar_queue = [ready_req, pending_req] + mgr.grammar_queue.join([ready_req, pending_req]) result = mgr.get_ready_grammar_requests() self.assertEqual(len(result), 1) self.assertIs(result[0], ready_req) self.assertEqual(len(mgr.grammar_queue), 1) - self.assertIs(mgr.grammar_queue[0], pending_req) + self.assertIs(list(mgr.grammar_queue)[0], pending_req) def test_empty_queue(self): """get_ready_grammar_requests on empty queue should return empty list.""" @@ -705,7 +705,7 @@ def test_progressive_timeout(self): req.grammar = future req.grammar_key = ("json", "slow") req.grammar_wait_ct = 2 # Already waited 2 iterations - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) # wait_ct increments to 3 (== max), should timeout result = mgr.get_ready_grammar_requests() @@ -725,7 +725,7 @@ def test_ready_future_applies_request_budget_without_polluting_cache(self): req = _make_req(json_schema="schema", custom_params={"thinking_budget": 4}) req.grammar = future req.grammar_key = ("json", "schema") - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) result = mgr.get_ready_grammar_requests() @@ -757,7 +757,7 @@ def test_multi_rank_sync_intersects_ready_unions_failed(self, mock_all_gather): req1.grammar_key = ("json", "s1") req1.grammar_wait_ct = 0 - mgr.grammar_queue = [req0, req1] + mgr.grammar_queue.join([req0, req1]) mgr.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = 100 # Simulate all_gather: rank 0 has {0} ready, rank 1 has {0,1} ready @@ -773,7 +773,7 @@ def fake_all_gather(output_list, _obj, group=None): # noqa: ARG001 self.assertIs(result[0], req0) # req1 stays in queue self.assertEqual(len(mgr.grammar_queue), 1) - self.assertIs(mgr.grammar_queue[0], req1) + self.assertIs(list(mgr.grammar_queue)[0], req1) @patch("sglang.srt.constrained.grammar_manager.torch.distributed.all_gather_object") def test_multi_rank_sync_unions_failed(self, mock_all_gather): @@ -788,7 +788,7 @@ def test_multi_rank_sync_unions_failed(self, mock_all_gather): req.grammar_key = ("json", "slow") req.grammar_wait_ct = 0 - mgr.grammar_queue = [req] + mgr.grammar_queue.join([req]) # Simulate: rank 0 has no ready and idx 0 failed, rank 1 has no ready/failed def fake_all_gather(output_list, _obj, group=None): # noqa: ARG001 @@ -1042,7 +1042,7 @@ def test_future_exception_creates_invalid_grammar(self): req.require_reasoning = True req.grammar = future req.grammar_key = ("json", '{"type": "object"}') - mgr.grammar_queue.append(req) + mgr.grammar_queue.join([req]) mgr.SGLANG_GRAMMAR_POLL_INTERVAL = 0.001 result = mgr.get_ready_grammar_requests() diff --git a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py index 0ff3ef20159e..5883bc5c2580 100644 --- a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py +++ b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py @@ -31,6 +31,10 @@ from sglang.srt.disaggregation.utils import is_unadmitted_reject from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.runtime_context import publish, reset_context from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs @@ -78,7 +82,7 @@ def _make_scheduler(): def _prefill_queue(sched): q = SimpleNamespace( scheduler=sched, - queue=[], + queue=_RequestList(ListKind.PD_BOOTSTRAP), create_sender=MagicMock(return_value=True), ) return q @@ -88,7 +92,7 @@ def _decode_queue(sched): q = SimpleNamespace( scheduler=sched, token_to_kv_pool_allocator=MagicMock(), - retracted_queue=[], + retracted_queue=_RequestList(ListKind.PD_RETRACTED), pending_reqs=[], _check_if_req_exceed_kv_capacity=MagicMock(return_value=False), _create_receiver_and_enqueue=MagicMock( @@ -142,7 +146,7 @@ def test_rejected_request_is_retired_at_each_pd_door(self): q = _prefill_queue(sched) _admit_prefill(q, req) q.create_sender.assert_not_called() - self.assertEqual(q.queue, []) + self.assertEqual(list(q.queue), []) else: q = _decode_queue(sched) _admit_decode(q, req) @@ -174,7 +178,7 @@ def test_reentering_prefill_request_is_not_retired(self): sched = _make_scheduler() q = _prefill_queue(sched) _admit_prefill(q, req) - self.assertEqual(q.queue, [req]) + self.assertEqual(list(q.queue), [req]) sched.retire_unadmitted_request.assert_not_called() self.assertIsInstance(req.to_finish, FINISH_ABORT) @@ -212,7 +216,7 @@ def test_retracted_decode_request_is_not_retired(self): sched.retire_unadmitted_request.assert_not_called() self.assertIsInstance(req.to_finish, FINISH_ABORT) if flag == "is_retracted": - self.assertEqual(q.retracted_queue, [req]) + self.assertEqual(list(q.retracted_queue), [req]) else: q._create_receiver_and_enqueue.assert_called_once() receiver = ( @@ -245,7 +249,7 @@ def test_preempted_reentry_is_not_retired(self): if door == "prefill": q = _prefill_queue(sched) _admit_prefill(q, req) - self.assertEqual(q.queue, [req]) + self.assertEqual(list(q.queue), [req]) else: q = _decode_queue(sched) _admit_decode(q, req) @@ -262,7 +266,7 @@ def test_valid_request_still_admitted(self): if door == "prefill": q = _prefill_queue(sched) _admit_prefill(q, req) - self.assertEqual(q.queue, [req]) + self.assertEqual(list(q.queue), [req]) else: q = _decode_queue(sched) _admit_decode(q, req) diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index d149a8f6be60..bb3f4c8bc2b0 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -69,7 +69,7 @@ def __init__(self): ) self.spec_algorithm = SimpleNamespace(is_eagle=lambda: False) self.tree_cache = Mock() - self.disagg_prefill_transfer_queue = [] + self.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) self.disagg_prefill_pending_chunk_rids = {"aborted-prefill"} self.send_kv_chunk = Mock() self.output_streamer = Mock() @@ -213,6 +213,7 @@ def ready(): ) req.disagg_kv_sender.mark_prefill_complete.side_effect = ready + scheduler.prefill_launched_reqs.join([req]) scheduler.process_batch_result_disagg_prefill(batch, result) self.assertEqual(observed, [([11], [1.0, 2.0], [5, 6], True)]) self.assertTrue(req.pending_bootstrap) @@ -225,7 +226,7 @@ def test_aborted_final_result_never_publishes_ready(self, release): scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) self.assertTrue(req.finished()) self.assertFalse(req.kv.holds_kv) - self.assertEqual(scheduler.disagg_prefill_transfer_queue, []) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), []) release.assert_called_once_with(req, scheduler.tree_cache, checkpoint=False) req.disagg_kv_sender.mark_prefill_complete.assert_not_called() @@ -235,11 +236,11 @@ def test_discarded_attempt_never_publishes_ready(self, force_retry): req = _Req() req.to_finish = None retried = [] - scheduler.optimistic_release_and_requeue = lambda r: retried.append(r) + scheduler.optimistic_release_and_requeue = lambda r, src: retried.append(r) scheduler.process_batch_result_disagg_prefill(_batch(req), _result()) self.assertEqual(retried, [req]) self.assertEqual(req.output_ids, []) - self.assertEqual(scheduler.disagg_prefill_transfer_queue, []) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), []) req.disagg_kv_sender.mark_prefill_complete.assert_not_called() diff --git a/test/registered/unit/disaggregation/test_prefill_complete_integration.py b/test/registered/unit/disaggregation/test_prefill_complete_integration.py index 20068515d768..3da698f7dbbf 100644 --- a/test/registered/unit/disaggregation/test_prefill_complete_integration.py +++ b/test/registered/unit/disaggregation/test_prefill_complete_integration.py @@ -5,7 +5,7 @@ from collections import defaultdict, deque from concurrent.futures import ThreadPoolExecutor, TimeoutError from types import MethodType, SimpleNamespace -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, Mock, call, patch import torch @@ -38,6 +38,7 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_request_lists import FakeReq from sglang.test.test_utils import CustomTestCase, enter_scope register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -332,7 +333,7 @@ def test_exhausted_optimistic_attempt_follows_allocation_policy(self): ) override.install() self.addCleanup(override.restore) - req = SimpleNamespace( + req = FakeReq( rid="retry", cache_request_handle=CacheRequestHandle(rid="retry", attempt_id=0), prefill_attempt_count=1, @@ -351,22 +352,27 @@ def test_exhausted_optimistic_attempt_follows_allocation_policy(self): scheduler = SimpleNamespace( tree_cache=Mock(supports_mamba=Mock(return_value=False)), clear_pending_chunk_send=Mock(), - prefill_launched_reqs=Mock(), + prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), waiting_queue=[], - disagg_prefill_bootstrap_queue=SimpleNamespace(queue=[]), + disagg_prefill_bootstrap_queue=SimpleNamespace( + queue=_RequestList(ListKind.PD_BOOTSTRAP) + ), metrics_reporter=SimpleNamespace(enable_metrics=False), processed_tokens_counter=0, req_to_metadata_buffer_idx_allocator=Mock(), ) + launched = scheduler.prefill_launched_reqs with ( patch("sglang.srt.disaggregation.prefill.checkpoint_kv_cache"), patch("sglang.srt.disaggregation.prefill.release_kv_cache"), ): + launched.join([req]) SchedulerDisaggregationPrefillMixin.optimistic_release_and_requeue( - scheduler, req + scheduler, req, launched ) self.assertEqual(scheduler.waiting_queue, [req]) - self.assertEqual(scheduler.disagg_prefill_bootstrap_queue.queue, []) + self.assertEqual(list(launched), []) + self.assertEqual(list(scheduler.disagg_prefill_bootstrap_queue.queue), []) self.assertEqual(req.prefill_attempt_count, 2) self.assertEqual(list(req.output_ids), []) self.assertTrue(req.pending_bootstrap) @@ -378,11 +384,15 @@ def test_exhausted_optimistic_attempt_follows_allocation_policy(self): disaggregation_decode_allocation_policy="early" ): scheduler.waiting_queue.clear() + launched.join([req]) SchedulerDisaggregationPrefillMixin.optimistic_release_and_requeue( - scheduler, req + scheduler, req, launched ) self.assertEqual(scheduler.waiting_queue, []) - self.assertEqual(scheduler.disagg_prefill_bootstrap_queue.queue, [req]) + self.assertEqual(list(launched), []) + self.assertEqual( + list(scheduler.disagg_prefill_bootstrap_queue.queue), [req] + ) self.assertEqual(req.metadata_buffer_index, -1) scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with( 0 @@ -436,7 +446,7 @@ class _Req: patch("sglang.srt.disaggregation.prefill.release_kv_cache"), ): SchedulerDisaggregationPrefillMixin.optimistic_release_and_requeue( - scheduler, req + scheduler, req, scheduler.prefill_launched_reqs ) self.assertEqual(scheduler.waiting_queue, [req]) @@ -486,14 +496,12 @@ def test_chunked_prefill_does_not_yield_for_decode_admission(self): else: self.assertIsNone(scheduler.chunked_req) # With a step in flight, it is requeued once processed. + launched = scheduler.prefill_launched_reqs self.assertEqual( - scheduler.optimistic_release_and_requeue.call_count, - 0 if inflight_steps else 1, - ) - self.assertEqual( - list(scheduler.prefill_launched_reqs), - [req] if inflight_steps else [], + scheduler.optimistic_release_and_requeue.call_args_list, + [] if inflight_steps else [call(req, launched)], ) + self.assertEqual(list(launched), [req]) finally: override.restore() @@ -525,6 +533,7 @@ def test_yielded_request_requeues_after_its_chunk_once_bootstrap_completes(self) ), spec_algorithm=SimpleNamespace(is_eagle=lambda: False), optimistic_release_and_requeue=Mock(), + prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()), maybe_send_health_check_signal=Mock(), ) @@ -533,13 +542,16 @@ def test_yielded_request_requeues_after_its_chunk_once_bootstrap_completes(self) scheduler, batch, result ) - scheduler.optimistic_release_and_requeue.assert_called_once_with(req) + scheduler.optimistic_release_and_requeue.assert_called_once_with( + req, scheduler.prefill_launched_reqs + ) def inflight_scheduler(self, reqs): scheduler = SchedulerDisaggregationPrefillMixin() scheduler.tree_cache = Mock() scheduler.scheduler_stage_metrics = None - scheduler.disagg_prefill_transfer_queue = list(reqs) + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) + scheduler.disagg_prefill_transfer_queue.join(reqs) scheduler.attn_cp_cpu_group = scheduler.attn_tp_cpu_group = object() scheduler.output_streamer = Mock() scheduler.send_kv_chunk = Mock() @@ -563,7 +575,7 @@ def test_completed_prefill_parks_until_normal_pointer_handshake(self): self.assertEqual(receiver.poll(), KVPoll.WaitingForInput) # Decode may reserve KV now, but prefill still has no destination indices. self.assertEqual(sender.poll(), KVPoll.Bootstrapping) - req = SimpleNamespace( + req = FakeReq( rid="parked", disagg_kv_sender=sender, pending_bootstrap=True, @@ -582,7 +594,7 @@ def finalize(r): finalize_bootstrap=Mock(side_effect=finalize) ) self.assertEqual(scheduler.process_disagg_prefill_transfer_queue(), []) - self.assertEqual(scheduler.disagg_prefill_transfer_queue, [req]) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), [req]) scheduler.send_kv_chunk.assert_not_called() # This is the existing metadata-received transition, not readiness. self.prefill.update_status(1, KVPoll.WaitingForInput) @@ -600,20 +612,24 @@ def test_parked_prefill_reclaims_only_one_slot(self): pool.alloc(reqs) scheduler.waiting_queue = [] scheduler.disagg_prefill_bootstrap_queue = SimpleNamespace(queue=[]) - scheduler.optimistic_release_and_requeue = Mock(side_effect=pool.free) + scheduler.optimistic_release_and_requeue = Mock( + side_effect=lambda req, src: pool.free(req) + ) with get_context().override_server_args( disaggregation_decode_allocation_policy="early" ): scheduler.process_disagg_prefill_transfer_queue() - self.assertEqual(scheduler.disagg_prefill_transfer_queue, reqs) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), reqs) scheduler.waiting_queue.append(object()) scheduler.process_disagg_prefill_transfer_queue() - self.assertEqual(scheduler.disagg_prefill_transfer_queue, reqs[1:]) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), reqs[1:]) self.assertEqual(pool.available_size(), 1) scheduler.process_disagg_prefill_transfer_queue() - scheduler.optimistic_release_and_requeue.assert_called_once_with(reqs[0]) + scheduler.optimistic_release_and_requeue.assert_called_once_with( + reqs[0], scheduler.disagg_prefill_transfer_queue + ) def test_success_is_sent_once_after_all_sender_state_is_cleared(self): sender, receiver = self.sender(), self.receiver() diff --git a/test/registered/unit/disaggregation/test_prefill_result_polling.py b/test/registered/unit/disaggregation/test_prefill_result_polling.py index a00e133b6997..91b1de000f77 100644 --- a/test/registered/unit/disaggregation/test_prefill_result_polling.py +++ b/test/registered/unit/disaggregation/test_prefill_result_polling.py @@ -121,6 +121,12 @@ def _input(batch): ) +def _sending(reqs=()): + sending = _RequestList(ListKind.PD_SENDING) + sending.join(reqs) + return sending + + class _Scheduler(SchedulerDisaggregationPrefillMixin): """Run real scheduling and result handling with controlled intake and GPU events.""" @@ -229,7 +235,7 @@ def __init__( self._record_step_counters = Mock() self._maybe_clear_mm_inputs = Mock() self.maybe_send_health_check_signal = Mock() - self.disagg_prefill_transfer_queue = [] + self.disagg_prefill_transfer_queue = _sending() self.send_kv_chunk = Mock() self.process_pending_chunked_abort = Mock() self.dp_attn_adapter = SimpleNamespace( @@ -320,6 +326,9 @@ def run_batch(self, batch): assert len(self.result_queue) < 2, "Overwrote a live forward's keep-alive slot" self.actions.append((self.iteration, "launch", batch.forward_iter)) batch.add_inflight_step() + self.prefill_launched_reqs.join( + req for req in batch.reqs if req is not batch.chunked_req + ) result = self.results[batch.forward_iter] if result.delay_sample_func is not None: @@ -472,7 +481,7 @@ def process_disagg_prefill_transfer_queue(self): ) scheduler.polls = 0 scheduler.tp_cpu_group = dist.group.WORLD - scheduler.disagg_prefill_transfer_queue = ["transferring"] + scheduler.disagg_prefill_transfer_queue = _sending(["transferring"]) with envs.SGLANG_DISAGG_PREFILL_POLL_TRANSFERS_WHILE_WAITING.override(True): _run_loop( scheduler, @@ -713,7 +722,7 @@ def process_disagg_prefill_transfer_queue(self): scheduler = TransferScheduler( [_batch(1)] + [None] * 3, {1: _result()}, {1: 4} ) - scheduler.disagg_prefill_transfer_queue = list(inflight) + scheduler.disagg_prefill_transfer_queue = _sending(inflight) with envs.SGLANG_DISAGG_PREFILL_POLL_TRANSFERS_WHILE_WAITING.override( poll_transfers ): @@ -763,7 +772,7 @@ def _inputs_for_step(self): pause_mode=event if pause else "in_place", ) scheduler.shutdown_at = None if pause else 2 - scheduler.disagg_prefill_transfer_queue = ["transferring"] + scheduler.disagg_prefill_transfer_queue = _sending(["transferring"]) with envs.SGLANG_DISAGG_PREFILL_POLL_TRANSFERS_WHILE_WAITING.override( True ): @@ -829,7 +838,7 @@ def test_transfer_wait_polling_needs_the_flag_and_a_supported_loop(self): {1: _result(_Event(ready_after_queries=3))}, {1: 4}, ) - scheduler.disagg_prefill_transfer_queue = ["transferring"] + scheduler.disagg_prefill_transfer_queue = _sending(["transferring"]) topology = {} if mode == "independent_dp": topology = {"dp_size": 2} @@ -943,7 +952,7 @@ def get_new_batch_prefill(self, running_batch): ), clear=Mock(), ) - scheduler.disagg_prefill_transfer_queue = [prior] + scheduler.disagg_prefill_transfer_queue = _sending([prior]) first.reqs[0].metadata_buffer_index = allocator.alloc() for req in (first.reqs[0], second.reqs[0]): req.pending_bootstrap = False @@ -1042,7 +1051,7 @@ def test_coarse_war_barrier_follows_previous_result_processing(self): {1: ready_at, 2: ready_at + 1}, ) scheduler.war_barrier_is_coarse = coarse - scheduler.disagg_prefill_transfer_queue = ["transferring"] + scheduler.disagg_prefill_transfer_queue = _sending(["transferring"]) with envs.SGLANG_DISAGG_PREFILL_POLL_TRANSFERS_WHILE_WAITING.override( polling == "transfers" ): @@ -1122,7 +1131,7 @@ def checkpoint_disagg_prefill(self, req): self.assertEqual(checkpoint_ends, [4, 8, 12]) self.assertEqual(req.inflight_steps, 0) self.assertEqual(list(req.output_ids), [11]) - self.assertEqual(scheduler.disagg_prefill_transfer_queue, [req]) + self.assertEqual(list(scheduler.disagg_prefill_transfer_queue), [req]) self.assertEqual(scheduler.result_processes, {4: 1, 8: 1, 12: 1}) for end, event in scheduler.copy_events.items(): self.assertEqual(event.synchronizations, 1) diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 9a872afddd39..73fdbf5abbef 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -19,6 +19,10 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.scheduler_pp_mixin import PPBatchMetadata from sglang.srt.managers.utils import GenerationBatchResult, get_logprob_from_pp_outputs from sglang.srt.model_executor.forward_batch_info import PPProxyTensors @@ -28,6 +32,7 @@ from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_request_lists import FakeReq from sglang.test.test_utils import published_topology register_cpu_ci(est_time=12, suite="base-a-test-cpu") @@ -513,7 +518,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): copy_done=copy_done, auxiliary_host_output=host_output, ) - req = SimpleNamespace( + req = FakeReq( output_ids=[], finished_len=None, to_finish=None, @@ -548,14 +553,15 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): batch_result_processor=processor, spec_algorithm=SimpleNamespace(is_eagle=lambda: False), tree_cache=object(), - disagg_prefill_transfer_queue=[], - prefill_launched_reqs=Mock(), + disagg_prefill_transfer_queue=_RequestList(ListKind.PD_SENDING), + prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), checkpoint_disagg_prefill=Mock(), send_kv_chunk=Mock(), metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()), maybe_send_health_check_signal=Mock(), ) + scheduler.prefill_launched_reqs.join([req]) SchedulerDisaggregationPrefillMixin.process_batch_result_disagg_prefill( scheduler, batch, diff --git a/test/registered/unit/managers/test_deferred_output.py b/test/registered/unit/managers/test_deferred_output.py index 4c666541dc28..30e3f33d52e4 100644 --- a/test/registered/unit/managers/test_deferred_output.py +++ b/test/registered/unit/managers/test_deferred_output.py @@ -61,7 +61,7 @@ def _scheduler(): scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None - scheduler.disagg_prefill_transfer_queue = None + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) scheduler.disagg_decode_prealloc_queue = None scheduler.disagg_decode_transfer_queue = None scheduler.ipc_channels = Mock() diff --git a/test/registered/unit/managers/test_multimodal_abort_cleanup.py b/test/registered/unit/managers/test_multimodal_abort_cleanup.py index f457aa34a93d..0dde32730e09 100644 --- a/test/registered/unit/managers/test_multimodal_abort_cleanup.py +++ b/test/registered/unit/managers/test_multimodal_abort_cleanup.py @@ -19,9 +19,6 @@ from sglang.srt.managers.req_types import MultimodalDataItem, MultimodalProcessorOutput from sglang.srt.managers.schedule_batch import MultimodalInputs, Req from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.deferred_responses import ( - DeferredResponses, -) from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.multimodal.modality import Modality from sglang.srt.multimodal.transport.cuda_ipc import CudaIpcTensorTransportProxy @@ -30,6 +27,7 @@ from sglang.srt.session.session_controller import Session from sglang.srt.utils import cuda_vmm_transport_utils as vmm from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_request_lists import empty_request_lists register_cpu_ci(est_time=12, suite="base-a-test-cpu") @@ -89,20 +87,15 @@ def _rank(rank): def _queued_scheduler(req): scheduler = object.__new__(Scheduler) - scheduler.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) + empty_request_lists(scheduler) scheduler.waiting_queue = [req] - scheduler.chunked_req = scheduler.mm_receiver = scheduler.dllm_config = None + scheduler.mm_receiver = None scheduler.disaggregation_mode = DisaggregationMode.NULL scheduler.tree_cache = MagicMock() scheduler.beam_coordinator = MagicMock() scheduler.ipc_channels = MagicMock() scheduler.rust_server = None scheduler.output_streamer = MagicMock() - scheduler.grammar_manager = MagicMock() scheduler.collect_inflight_reqs = lambda: set() return scheduler diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 84b1abfe4466..ba3fd5b74224 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -28,8 +28,9 @@ ReqKvInfo, ) from sglang.srt.managers.scheduler import Scheduler # noqa: E402 -from sglang.srt.managers.scheduler_components.deferred_responses import ( - DeferredResponses, +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, ) from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache # noqa: E402 @@ -38,6 +39,7 @@ from sglang.srt.runtime_context import get_context, publish, reset_context # noqa: E402 from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_request_lists import FakeReq, empty_request_lists from sglang.test.separate_buffer_allocator_double import ( bind_separate_buffer_capacity, ) @@ -238,7 +240,7 @@ def _pop_bootstrapping_swa_branch( processed_tokens_counter=0, ) queue.ensure_metadata_buffer = MagicMock(return_value=True) - req = SimpleNamespace( + req = FakeReq( rid="req", disagg_kv_sender=object(), prefill_attempt_count=0, @@ -247,7 +249,8 @@ def _pop_bootstrapping_swa_branch( kv=SimpleNamespace(cache_protected_len=2), time_stats=MagicMock(), ) - queue.queue = [req] + queue.queue = _RequestList(ListKind.PD_BOOTSTRAP) + queue.queue.join([req]) with ( patch( @@ -267,7 +270,7 @@ def test_write_through_pending_swa_branch_waits_for_bootstrap(self): queue, req, result = self._pop_bootstrapping_swa_branch("write_through") self.assertEqual(result, []) - self.assertEqual(queue.queue, [req]) + self.assertEqual(list(queue.queue), [req]) self.assertEqual(req.prefill_attempt_count, 0) queue.ensure_metadata_buffer.assert_not_called() @@ -277,7 +280,7 @@ def test_other_write_policies_stay_optimistic(self): queue, req, result = self._pop_bootstrapping_swa_branch(write_policy) self.assertEqual(result, [req]) - self.assertEqual(queue.queue, []) + self.assertEqual(list(queue.queue), []) self.assertEqual(req.prefill_attempt_count, 1) queue.ensure_metadata_buffer.assert_called_once_with(req) @@ -299,7 +302,7 @@ def can_prefill_before_bootstrap(self, req): "write_through", queue_cls=_Eager, is_retracted=True ) self.assertEqual(result, []) - self.assertEqual(queue.queue, [req]) + self.assertEqual(list(queue.queue), [req]) queue.ensure_metadata_buffer.assert_not_called() def test_unpublished_advance_preserves_swa_frontier_and_ownership(self): @@ -435,12 +438,7 @@ def test_waiting_abort_releases_without_insert(self): def test_scheduler_waiting_abort_delegates_release_once(self): scheduler = Scheduler.__new__(Scheduler) - scheduler.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) - scheduler.chunked_req = None + empty_request_lists(scheduler) scheduler.mm_receiver = None scheduler.waiting_queue = [MagicMock(rid="req")] scheduler.beam_coordinator = MagicMock() @@ -450,10 +448,6 @@ def test_scheduler_waiting_abort_delegates_release_once(self): scheduler.output_streamer = MagicMock() scheduler.disaggregation_mode = DisaggregationMode.PREFILL scheduler.release_aborted_prefill_waiting_req = MagicMock() - scheduler.dllm_config = None - scheduler.grammar_manager = MagicMock() - scheduler.disagg_prefill_bootstrap_queue = SimpleNamespace(queue=[]) - scheduler.disagg_prefill_transfer_queue = [] scheduler.collect_inflight_reqs = MagicMock(return_value=set()) scheduler.tree_cache = MagicMock() diff --git a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py index 019523537100..ff9e9c084c47 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py +++ b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py @@ -15,12 +15,9 @@ maybe_stub_sgl_kernel() -from sglang.srt.disaggregation.mode import DisaggregationMode # noqa: E402 from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req # noqa: E402 from sglang.srt.managers.scheduler import Scheduler # noqa: E402 -from sglang.srt.managers.scheduler_components.deferred_responses import ( # noqa: E402 - DeferredResponses, -) +from sglang.test.scheduler_request_lists import empty_request_lists # noqa: E402 register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -51,13 +48,11 @@ def finished(self): def _make_scheduler(pending_req, *, chunked_req, running_reqs) -> Scheduler: sched = Scheduler.__new__(Scheduler) + empty_request_lists(sched) sched.enable_continuous_input_polling = False sched.result_queue = deque() sched.chunked_req = chunked_req sched._pending_chunked_abort_req = pending_req - sched.waiting_queue = [] - sched.dllm_config = None - sched.grammar_manager = Mock() sched.disaggregation_mode = None sched.enable_hicache_storage = False sched.mm_receiver = None @@ -66,11 +61,6 @@ def _make_scheduler(pending_req, *, chunked_req, running_reqs) -> Scheduler: sched.last_batch = None sched.beam_coordinator = Mock() sched.output_streamer = Mock() - sched.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) return sched diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py index 42722c242442..8092ef19578b 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py +++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py @@ -67,7 +67,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None - scheduler.disagg_prefill_transfer_queue = None + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) scheduler.disagg_decode_prealloc_queue = None scheduler.disagg_decode_transfer_queue = None scheduler.hisparse_coordinator = None diff --git a/test/registered/unit/managers/test_scheduler_on_idle_load.py b/test/registered/unit/managers/test_scheduler_on_idle_load.py index cfb78cb2b057..d0f3f562848a 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/test/registered/unit/managers/test_scheduler_on_idle_load.py @@ -127,7 +127,7 @@ def _idle_scheduler(self, materializer): scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None - scheduler.disagg_prefill_transfer_queue = None + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) scheduler.disagg_decode_prealloc_queue = None scheduler.disagg_decode_transfer_queue = None scheduler.hisparse_coordinator = None diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 0e4d3a3bcfae..547bb8a27068 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -13,7 +13,7 @@ maybe_stub_sgl_kernel() -from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeTransferQueue from sglang.srt.disaggregation.decode_kvcache_offload_manager import ( DecodeKVCacheOffloadManager, ) @@ -131,7 +131,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.hisparse_coordinator = None scheduler.decode_offload_manager = None scheduler.disagg_prefill_bootstrap_queue = None - scheduler.disagg_prefill_transfer_queue = None + scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) scheduler.disagg_decode_prealloc_queue = None scheduler.disagg_decode_transfer_queue = None scheduler.dllm_config = None @@ -899,12 +899,18 @@ def test_pd_decode_abort_finishes_a_held_rebootstrap(self): scheduler = self._new_scheduler() scheduler.disaggregation_mode = DisaggregationMode.DECODE prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + prealloc_queue.scheduler = scheduler prealloc_queue.queue = [] - prealloc_queue.retracted_queue = [] - prealloc_queue.held_rebootstrap_reqs = [] + prealloc_queue.retracted_queue = _RequestList(ListKind.PD_RETRACTED) + prealloc_queue.held_rebootstrap_reqs = _RequestList( + ListKind.PD_HELD, abort=prealloc_queue._abort_held + ) prealloc_queue.add = MagicMock() scheduler.disagg_decode_prealloc_queue = prealloc_queue - scheduler.disagg_decode_transfer_queue = SimpleNamespace(queue=[]) + transfer_queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + transfer_queue.queue = [] + transfer_queue._deferred_releases = [] + scheduler.disagg_decode_transfer_queue = transfer_queue scheduler.deferred_responses = MagicMock() scheduler.mm_receiver = None scheduler.grammar_manager = MagicMock() @@ -919,7 +925,7 @@ def test_pd_decode_abort_finishes_a_held_rebootstrap(self): self.assertIsInstance(aborted.finished_reason, FINISH_ABORT) self.assertIsNone(kept.finished_reason) - self.assertEqual(prealloc_queue.held_rebootstrap_reqs, [kept]) + self.assertEqual(list(prealloc_queue.held_rebootstrap_reqs), [kept]) scheduler.tree_cache.finish.assert_called_once() ((abort, sent_req), _) = scheduler.send_abort_output.call_args self.assertEqual((abort.rid, sent_req), ("aborted", aborted)) diff --git a/test/registered/unit/managers/test_scheduler_running_reqs.py b/test/registered/unit/managers/test_scheduler_running_reqs.py index 71381297f17c..826aba928edb 100644 --- a/test/registered/unit/managers/test_scheduler_running_reqs.py +++ b/test/registered/unit/managers/test_scheduler_running_reqs.py @@ -25,9 +25,6 @@ ) from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.deferred_responses import ( - DeferredResponses, -) from sglang.srt.managers.scheduler_components.invariant_checker import ( SchedulerInvariantChecker, ) @@ -38,6 +35,7 @@ from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.scheduler_request_lists import empty_request_lists register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -94,34 +92,14 @@ def _invariant_checker(s): ) -def _empty_request_lists(s): - s.waiting_queue = [] - s.grammar_manager = SimpleNamespace(grammar_queue=[]) - s.chunked_req = None - s.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) - s.disagg_prefill_transfer_queue = None - s.disagg_prefill_bootstrap_queue = None - s.disagg_decode_prealloc_queue = None - s.disagg_decode_transfer_queue = None - s.dllm_manager = SimpleNamespace(waiting_queue=[]) - s.hisparse_coordinator = None - s.uocr_runtime = None - s.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) - - def _scheduler(): s = _Scheduler.__new__(_Scheduler) + empty_request_lists(s) + s.dllm_manager = SimpleNamespace(waiting_queue=[]) s.running_batch = ScheduleBatch(reqs=[]) - s.running_reqs = _RequestList(ListKind.DECODING) s.invariant_checker = _invariant_checker(s) s.disaggregation_mode = DisaggregationMode.NULL - s.dllm_config = None s.enable_hisparse = False - _empty_request_lists(s) # run_batch up to the forward s.scheduler_stage_metrics = None s.metrics_reporter = MagicMock() @@ -320,14 +298,12 @@ def test_ngram_grammar_step_launches_no_finished_row(self): running.spec_algorithm = SpeculativeAlgorithm.NGRAM running.has_grammar = True s = Scheduler.__new__(Scheduler) + empty_request_lists(s) s.running_batch = running s.last_batch = None - s.running_reqs = _RequestList(ListKind.DECODING) s.running_reqs.join(running.reqs) s.disaggregation_mode = DisaggregationMode.NULL - s.dllm_config = None s.enable_hisparse = False - _empty_request_lists(s) s.invariant_checker = _invariant_checker(s) s._confidence_budget_prepare = None # Event loop diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 237450b97136..3931758a9e57 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -27,13 +27,11 @@ from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.deferred_responses import ( - DeferredResponses, -) from sglang.srt.mem_cache.base_prefix_cache import ( CacheRequestHandle, CacheRequestOutcome, ) +from sglang.test.scheduler_request_lists import empty_request_lists register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -77,11 +75,7 @@ def _batch(reqs): def _scheduler(waiting_queue, running_reqs=(), last_batch_reqs=()): s = Scheduler.__new__(Scheduler) - s.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) + empty_request_lists(s) s.enable_continuous_input_polling = False s.result_queue = deque() s.waiting_queue = waiting_queue diff --git a/test/registered/unit/mem_cache/test_decode_retraction_backup.py b/test/registered/unit/mem_cache/test_decode_retraction_backup.py index 88db379b3324..cf20d02a5550 100644 --- a/test/registered/unit/mem_cache/test_decode_retraction_backup.py +++ b/test/registered/unit/mem_cache/test_decode_retraction_backup.py @@ -17,6 +17,10 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import ( SchedulerPoolStatsObserver, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.managers.utils import allocate_distinct_stream from sglang.srt.mem_cache.allocator import ( PagedTokenToKVPoolAllocator, @@ -371,7 +375,7 @@ def _receive_queue(self, env): transfer_queue=SimpleNamespace(queue=[], enable_staging=False), queue=[], pending_reqs=[], - retracted_queue=[], + retracted_queue=_RequestList(ListKind.PD_RETRACTED), _prefill_dp_rank_queries={}, _num_published_destinations=0, num_reserved_decode_tokens=0, @@ -769,7 +773,8 @@ def pre_alloc(resumed): (resumed.kv.req_pool_idx, slice(0, self.num_tokens)), reused ) - queue.retracted_queue = [req] + queue.retracted_queue = _RequestList(ListKind.PD_RETRACTED) + queue.retracted_queue.join([req]) with ( patch.object(queue, "_uses_swa_reservation", return_value=False), patch.object(queue, "_allocatable_token_budgets", return_value=1 << 20), diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index 92c89693651f..730fdc0f98dd 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -11,6 +11,7 @@ from sglang.srt.dllm.algorithm.low_confidence import LowConfidence from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.mixin.scheduler import DllmManager, SchedulerDllmMixin +from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.schedule_batch import Req, ReqKvInfo, ScheduleBatch from sglang.srt.mem_cache.allocation import alloc_for_extend from sglang.srt.mem_cache.base_prefix_cache import InsertParams @@ -229,21 +230,22 @@ def test_alloc_for_extend_rejects_partial_retained_block_reuse(self): with self.assertRaisesRegex(RuntimeError, "full block"): alloc_for_extend(batch) - def test_dllm_manager_pop_aborted_reqs_removes_waiting_and_staging(self): - manager = DllmManager(SimpleNamespace(max_running_requests=4)) - waiting = _make_req("abort-waiting", [1], self.block_size) - staging = _make_req("abort-staging", [2], self.block_size) - keep = _make_req("keep", [3], self.block_size) - manager.waiting_queue = [waiting, keep] - manager.staging_queue = [staging, waiting] + def test_dllm_abort_removes_waiting_and_staging(self): + released = [] + manager = DllmManager( + SimpleNamespace(max_running_requests=4), release_aborted=released.append + ) + staged, waiting, keep = ( + Mock(rid=rid) for rid in ("abort-staged", "abort-waiting", "keep") + ) + manager.waiting_queue.join([staged, waiting, keep]) + manager.staging_queue = [staged, keep] - aborted = manager.pop_aborted_reqs(False, "abort") + manager.waiting_queue.abort(AbortReq(rid="abort")) - self.assertEqual( - [req.rid for req in aborted], ["abort-waiting", "abort-staging"] - ) - self.assertEqual(manager.waiting_queue, [keep]) - self.assertEqual(manager.staging_queue, []) + self.assertEqual(released, [staged, waiting]) + self.assertEqual(list(manager.waiting_queue), [keep]) + self.assertEqual(manager.staging_queue, [keep]) class TestDllmFdfoResolvedBlockKeepsRow(unittest.TestCase): diff --git a/test/registered/unit/mem_cache/test_mamba_capacity.py b/test/registered/unit/mem_cache/test_mamba_capacity.py index f5f835d590a0..7efaf2458d08 100644 --- a/test/registered/unit/mem_cache/test_mamba_capacity.py +++ b/test/registered/unit/mem_cache/test_mamba_capacity.py @@ -7,6 +7,10 @@ DecodeReqToTokenPool, HybridMambaDecodeReqToTokenPool, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.mem_cache.allocation import alloc_req_slots, ensure_mamba_capacity from sglang.srt.mem_cache.allocation_sizing import get_mamba_tracking_slots from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator @@ -14,6 +18,7 @@ from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.static_kv_storage import StaticKVStorage from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_request_lists import FakeReq register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -226,7 +231,7 @@ def test_tracking_buffer_settings(self): def test_pd_preallocation_leaves_blocked_request_queued(self): pool, held = make_pool(1) cache = make_cache(pool, held, 0) - req = SimpleNamespace( + req = FakeReq( rid="blocked", origin_input_ids=[1], output_ids=[], @@ -241,7 +246,8 @@ def test_pd_preallocation_leaves_blocked_request_queued(self): queue.pp_size = 1 queue.queue = [decode_req] queue.pending_reqs = [] - queue.retracted_queue = [req] + queue.retracted_queue = _RequestList(ListKind.PD_RETRACTED) + queue.retracted_queue.join([req]) queue.req_to_token_pool = pool queue.token_to_kv_pool_allocator = SimpleNamespace(state=pool.state) queue.req_to_metadata_buffer_idx_allocator = SimpleNamespace( @@ -287,7 +293,7 @@ def test_pd_preallocation_leaves_blocked_request_queued(self): ), ): self.assertEqual(queue.resume_retracted_reqs(), []) - self.assertEqual(queue.retracted_queue, [req]) + self.assertEqual(list(queue.retracted_queue), [req]) self.assertTrue(req.is_retracted) queue._pre_alloc.assert_not_called() From 12a039bf3fb9041dd1448a50f5234fa689730c77 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:11:27 -0800 Subject: [PATCH 04/12] add a test that every field holding requests is a request list --- .../scheduler_components/test_request_list.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 test/registered/unit/managers/scheduler_components/test_request_list.py diff --git a/test/registered/unit/managers/scheduler_components/test_request_list.py b/test/registered/unit/managers/scheduler_components/test_request_list.py new file mode 100644 index 000000000000..ad2e98977667 --- /dev/null +++ b/test/registered/unit/managers/scheduler_components/test_request_list.py @@ -0,0 +1,82 @@ +import ast +import inspect +import re +import textwrap +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.constrained.grammar_manager import GrammarManager +from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeTransferQueue +from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue +from sglang.srt.dllm.mixin.scheduler import DllmManager +from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.deferred_responses import ( + DeferredResponses, +) +from sglang.srt.models.unlimited_ocr_v2_support.scheduler_runtime import ( + UocrSchedulerRuntime, +) + +register_cpu_ci(est_time=12, suite="base-a-test-cpu") + +# The classes that own the request lists, with their mixins. +_OWNERS = ( + Scheduler, + GrammarManager, + PrefillBootstrapQueue, + DecodePreallocQueue, + DecodeTransferQueue, + DllmManager, + HiSparseCoordinator, + UocrSchedulerRuntime, + DeferredResponses, +) + +# Fields annotated as holding requests that are not a `RequestList`. Only +# shrinks: a list that migrates leaves it. +_NOT_YET_LISTS = { + "Scheduler.waiting_queue", + "DecodePreallocQueue.queue", + # Shares its entries with `queue`. + "DecodePreallocQueue.pending_reqs", + "DecodeTransferQueue.queue", + "DecodeTransferQueue._deferred_releases", + # A subset of `waiting_queue`. + "DllmManager.staging_queue", + "DeferredResponses.held", +} + +_HOLDS_REQUESTS = re.compile(r"\b(Req|DecodeRequest)\b") + + +def _fields_holding_requests(owner) -> set: + fields = set() + for klass in owner.__mro__: + if not klass.__module__.startswith("sglang."): + continue + tree = ast.parse(textwrap.dedent(inspect.getsource(klass))) + for node in ast.walk(tree): + if ( + isinstance(node, ast.AnnAssign) + and isinstance(node.target, ast.Attribute) + and isinstance(node.target.value, ast.Name) + and node.target.value.id == "self" + and _HOLDS_REQUESTS.search(ast.unparse(node.annotation)) + ): + fields.add(f"{owner.__name__}.{node.target.attr}") + return fields + + +class TestRequestHoldersAreLists(CustomTestCase): + def test_every_field_holding_requests_is_a_request_list(self): + fields = set().union(*map(_fields_holding_requests, _OWNERS)) + self.assertEqual(fields, _NOT_YET_LISTS) + + +if __name__ == "__main__": + unittest.main() From 7311ececab6ff43c45d2e30dffcdb8473927db3e Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:46:19 -0800 Subject: [PATCH 05/12] number list kinds with enum auto --- .../scheduler_components/request_list.py | 37 +++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index d0e4db10272a..649a4738055b 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -18,28 +18,27 @@ from sglang.srt.managers.schedule_batch import Req -class ListKind(enum.IntEnum): - """The lists that hold the live requests. Values below 16 are the core's own - lists; the lists of other owners take values from 16 up.""" +class ListKind(enum.Enum): + """The lists that hold the live requests.""" - WAITING = 1 + WAITING = enum.auto() # The running requests. - DECODING = 2 - CHUNKED = 3 - PD_BOOTSTRAP = 16 + DECODING = enum.auto() + CHUNKED = enum.auto() + PD_BOOTSTRAP = enum.auto() # Prefill requests out of the chunked slot whose launched steps are not done. - PD_PREFILLING = 17 - PD_SENDING = 18 - PD_PREALLOC = 24 - PD_RETRACTED = 25 - PD_HELD = 26 - PD_RECEIVING = 27 - PD_RELEASING = 29 - HISPARSE_STAGING = 32 - GRAMMAR = 36 - DLLM = 40 - UOCR_PARKED = 48 - DEFERRED = 56 + PD_PREFILLING = enum.auto() + PD_SENDING = enum.auto() + PD_PREALLOC = enum.auto() + PD_RETRACTED = enum.auto() + PD_HELD = enum.auto() + PD_RECEIVING = enum.auto() + PD_RELEASING = enum.auto() + HISPARSE_STAGING = enum.auto() + GRAMMAR = enum.auto() + DLLM = enum.auto() + UOCR_PARKED = enum.auto() + DEFERRED = enum.auto() class RequestList(ABC): From 37ab0cf0607afa124974ca48620477b1916054ab Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:56:15 -0800 Subject: [PATCH 06/12] name list abort callbacks _abort or _abort_ --- python/sglang/srt/constrained/grammar_manager.py | 4 ++-- python/sglang/srt/disaggregation/prefill.py | 6 +++--- python/sglang/srt/managers/scheduler.py | 2 +- .../unit/constrained/test_grammar_manager.py | 14 +++++++------- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 237e7f7db45d..cbadb18d9bfa 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -36,7 +36,7 @@ class GrammarManager: def __init__(self, scheduler: Scheduler): self.scheduler = scheduler self.server_args = scheduler.server_args - self.grammar_queue = _RequestList(ListKind.GRAMMAR, abort=self.abort_requests) + 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, @@ -115,7 +115,7 @@ def _pp_sync_ready_failed( ) return data - def abort_requests(self, recv_req: AbortReq): + def _abort(self, recv_req: AbortReq) -> None: # Abort method 2: call `set_finish_with_abort` # The request will still run one prefill forward pass. # In this case, we change the input_ids to be only one token to make this prefill cheap. diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 4722a7caf29d..365fa6adeb97 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -213,7 +213,7 @@ def __init__( self.pp_size = parallel.pp_size self.gpu_id = gpu_id self.bootstrap_port = bootstrap_port - self.queue = _RequestList(ListKind.PD_BOOTSTRAP, abort=self.abort_requests) + self.queue = _RequestList(ListKind.PD_BOOTSTRAP, abort=self._abort) self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics # Each rank owns a subset of logical KV pages, so capacity widens by N. @@ -481,7 +481,7 @@ def extend(self, reqs: List[Req], num_kv_heads: int) -> None: for req in reqs: self.add(req, num_kv_heads) - def abort_requests(self, recv_req: AbortReq) -> None: + def _abort(self, recv_req: AbortReq) -> None: # Abort requests that have not yet been bootstrapped for req in self.queue: if recv_req.abort_all or req.rid.startswith(recv_req.rid): @@ -1218,7 +1218,7 @@ def advance_logprob_pt(i: int, req: Req) -> None: ) self.maybe_send_health_check_signal() - def abort_disagg_prefill_transfer_reqs(self: Scheduler, recv_req: AbortReq) -> None: + def _abort_prefill_transfer(self: Scheduler, recv_req: AbortReq) -> None: # Abort in-flight requests for req in self.disagg_prefill_transfer_queue: if recv_req.abort_all or req.rid.startswith(recv_req.rid): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8617d1535357..4a3890e3f3fd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1575,7 +1575,7 @@ def init_disaggregation(self): self.mm_receiver = None self.disagg_prefill_bootstrap_queue = None self.disagg_prefill_transfer_queue = _RequestList( - ListKind.PD_SENDING, abort=self.abort_disagg_prefill_transfer_reqs + ListKind.PD_SENDING, abort=self._abort_prefill_transfer ) self.disagg_decode_prealloc_queue = None self.disagg_decode_transfer_queue = None diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index 4183f76b29bf..c28179394d35 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -6,7 +6,7 @@ - process_req_with_grammar: dispatch by constraint type (json, regex, ebnf, structural_tag), no-constraint requests, no-backend error, cache hits, cached invalid grammar abort -- abort_requests: single abort, abort all, future cancellation +- abort: single abort, abort all, future cancellation - get_ready_grammar_requests: future completion, invalid grammar handling, timeout with max poll iterations, aborted request handling, queue cleanup @@ -469,7 +469,7 @@ def setUp(self): reset_context() self.addCleanup(reset_context) - """Test abort_requests handling.""" + """Test grammar queue abort handling.""" def _make_mgr_with_queue(self): scheduler = _make_scheduler() @@ -491,7 +491,7 @@ def test_abort_by_rid_prefix(self): abort_req.abort_all = False abort_req.rid = "req-123" - mgr.abort_requests(abort_req) + mgr.grammar_queue.abort(abort_req) future.cancel.assert_called_once() req.set_finish_with_abort.assert_called_once() @@ -505,7 +505,7 @@ def test_abort_non_matching_rid(self): abort_req.abort_all = False abort_req.rid = "req-123" - mgr.abort_requests(abort_req) + mgr.grammar_queue.abort(abort_req) req.set_finish_with_abort.assert_not_called() def test_abort_all(self): @@ -521,7 +521,7 @@ def test_abort_all(self): abort_req.abort_all = True abort_req.rid = "" - mgr.abort_requests(abort_req) + mgr.grammar_queue.abort(abort_req) for req in reqs: req.set_finish_with_abort.assert_called_once() @@ -531,7 +531,7 @@ def test_abort_empty_queue(self): abort_req = MagicMock() abort_req.abort_all = True abort_req.rid = "" - mgr.abort_requests(abort_req) # Should not raise + mgr.grammar_queue.abort(abort_req) # Should not raise def test_abort_prefix_match(self): """rid.startswith means prefix matching, not exact matching.""" @@ -544,7 +544,7 @@ def test_abort_prefix_match(self): abort_req.abort_all = False abort_req.rid = "req-123" - mgr.abort_requests(abort_req) + mgr.grammar_queue.abort(abort_req) req.set_finish_with_abort.assert_called_once() From 8f3c8a2d4714452b14a11b0c36c5519a0f9d4335 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 20:58:43 -0800 Subject: [PATCH 07/12] rename scheduler_request_lists test helper to scheduler_test_utils --- .../{scheduler_request_lists.py => scheduler_test_utils.py} | 0 .../unit/disaggregation/test_prefill_complete_integration.py | 2 +- test/registered/unit/managers/test_auxiliary_output.py | 2 +- test/registered/unit/managers/test_multimodal_abort_cleanup.py | 2 +- .../unit/managers/test_priority_scheduling_disaggregation.py | 2 +- .../unit/managers/test_scheduler_chunked_abort_race.py | 2 +- test/registered/unit/managers/test_scheduler_running_reqs.py | 2 +- test/registered/unit/managers/test_scheduler_timeouts.py | 2 +- test/registered/unit/mem_cache/test_mamba_capacity.py | 2 +- 9 files changed, 8 insertions(+), 8 deletions(-) rename python/sglang/test/{scheduler_request_lists.py => scheduler_test_utils.py} (100%) diff --git a/python/sglang/test/scheduler_request_lists.py b/python/sglang/test/scheduler_test_utils.py similarity index 100% rename from python/sglang/test/scheduler_request_lists.py rename to python/sglang/test/scheduler_test_utils.py diff --git a/test/registered/unit/disaggregation/test_prefill_complete_integration.py b/test/registered/unit/disaggregation/test_prefill_complete_integration.py index 3da698f7dbbf..db0a1b607181 100644 --- a/test/registered/unit/disaggregation/test_prefill_complete_integration.py +++ b/test/registered/unit/disaggregation/test_prefill_complete_integration.py @@ -38,7 +38,7 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_request_lists import FakeReq +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase, enter_scope register_cpu_ci(est_time=11, suite="base-a-test-cpu") diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 73fdbf5abbef..77e05458ef0e 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -32,7 +32,7 @@ from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_request_lists import FakeReq +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import published_topology register_cpu_ci(est_time=12, suite="base-a-test-cpu") diff --git a/test/registered/unit/managers/test_multimodal_abort_cleanup.py b/test/registered/unit/managers/test_multimodal_abort_cleanup.py index 0dde32730e09..7655c2437dc1 100644 --- a/test/registered/unit/managers/test_multimodal_abort_cleanup.py +++ b/test/registered/unit/managers/test_multimodal_abort_cleanup.py @@ -27,7 +27,7 @@ from sglang.srt.session.session_controller import Session from sglang.srt.utils import cuda_vmm_transport_utils as vmm from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_request_lists import empty_request_lists +from sglang.test.scheduler_test_utils import empty_request_lists register_cpu_ci(est_time=12, suite="base-a-test-cpu") diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index ba3fd5b74224..dd5a6d8f91c2 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -39,7 +39,7 @@ from sglang.srt.runtime_context import get_context, publish, reset_context # noqa: E402 from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_request_lists import FakeReq, empty_request_lists +from sglang.test.scheduler_test_utils import FakeReq, empty_request_lists from sglang.test.separate_buffer_allocator_double import ( bind_separate_buffer_capacity, ) diff --git a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py index ff9e9c084c47..f528f9c4f5c8 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py +++ b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py @@ -17,7 +17,7 @@ from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req # noqa: E402 from sglang.srt.managers.scheduler import Scheduler # noqa: E402 -from sglang.test.scheduler_request_lists import empty_request_lists # noqa: E402 +from sglang.test.scheduler_test_utils import empty_request_lists # noqa: E402 register_cpu_ci(est_time=11, suite="base-a-test-cpu") diff --git a/test/registered/unit/managers/test_scheduler_running_reqs.py b/test/registered/unit/managers/test_scheduler_running_reqs.py index 826aba928edb..1ac9a3780638 100644 --- a/test/registered/unit/managers/test_scheduler_running_reqs.py +++ b/test/registered/unit/managers/test_scheduler_running_reqs.py @@ -35,7 +35,7 @@ from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.test.scheduler_request_lists import empty_request_lists +from sglang.test.scheduler_test_utils import empty_request_lists register_cpu_ci(est_time=10, suite="base-a-test-cpu") diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 3931758a9e57..37058db38891 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -31,7 +31,7 @@ CacheRequestHandle, CacheRequestOutcome, ) -from sglang.test.scheduler_request_lists import empty_request_lists +from sglang.test.scheduler_test_utils import empty_request_lists register_cpu_ci(est_time=11, suite="base-a-test-cpu") diff --git a/test/registered/unit/mem_cache/test_mamba_capacity.py b/test/registered/unit/mem_cache/test_mamba_capacity.py index 7efaf2458d08..290ef99737ba 100644 --- a/test/registered/unit/mem_cache/test_mamba_capacity.py +++ b/test/registered/unit/mem_cache/test_mamba_capacity.py @@ -18,7 +18,7 @@ from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.static_kv_storage import StaticKVStorage from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_request_lists import FakeReq +from sglang.test.scheduler_test_utils import FakeReq register_cpu_ci(est_time=10, suite="base-a-test-cpu") From 3c7e06fbcc98750334d497c7d8829c7125c2138c Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 21:03:43 -0800 Subject: [PATCH 08/12] add RequestList.clear; take held rebootstrap requests with take --- python/sglang/srt/disaggregation/decode.py | 5 ++--- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/disaggregation/role_switch.py | 3 +-- .../srt/managers/scheduler_components/request_list.py | 7 +++++++ 4 files changed, 11 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 401de66a29a3..d3f40669df65 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1025,8 +1025,7 @@ def hold_rebootstrap(self, req: Req) -> None: def enqueue_held_rebootstrap(self) -> None: """Enqueue all staged rebootstrap requests when generation resumes.""" - held = list(self.held_rebootstrap_reqs) - self.held_rebootstrap_reqs.leave(held) + held = self.held_rebootstrap_reqs.take(lambda req: True) for req in held: self.add(req, is_rebootstrap=True) @@ -1112,7 +1111,7 @@ def release_memory_occupation(self): self.tree_cache, get_disagg().disaggregation_decode_retraction_backup, ) - self.retracted_queue.leave(list(self.retracted_queue)) + self.retracted_queue.clear() if hasattr(self.kv_manager, "deregister_buffer_to_engine"): self.kv_manager.deregister_buffer_to_engine() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 365fa6adeb97..1132f0fece31 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -586,7 +586,7 @@ def pop_bootstrapped(self) -> List[Req]: return bootstrapped_reqs def release_memory_occupation(self): - self.queue.leave(list(self.queue)) + self.queue.clear() if hasattr(self.kv_manager, "deregister_buffer_to_engine"): self.kv_manager.deregister_buffer_to_engine() diff --git a/python/sglang/srt/disaggregation/role_switch.py b/python/sglang/srt/disaggregation/role_switch.py index 781873bf9b3b..83a432390596 100644 --- a/python/sglang/srt/disaggregation/role_switch.py +++ b/python/sglang/srt/disaggregation/role_switch.py @@ -247,8 +247,7 @@ def teardown_disaggregation(scheduler: Scheduler) -> None: if km is not None: km.teardown() scheduler.disagg_prefill_bootstrap_queue = None - transfer_queue = scheduler.disagg_prefill_transfer_queue - transfer_queue.leave(list(transfer_queue)) + scheduler.disagg_prefill_transfer_queue.clear() elif mode == DisaggregationMode.DECODE: q = getattr(scheduler, "disagg_decode_prealloc_queue", None) if q is not None: diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index 649a4738055b..e275e35477f0 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -64,6 +64,10 @@ def take(self, pred: Callable[[Req], bool]) -> List[Req]: self.leave(taken) return taken + def clear(self) -> None: + """Drop every request.""" + self.leave(list(self)) + def __len__(self) -> int: return sum(1 for _ in self) @@ -95,6 +99,9 @@ def leave(self, reqs: Iterable[Req]) -> None: def leave_finished(self) -> None: self._drop({req for req in self._reqs if req.finished()}) + def clear(self) -> None: + self._reqs = {} + def abort(self, recv_req: AbortReq) -> None: if self._abort is not None: self._abort(recv_req) From e379320591f01b6e30c76b365348b45c12d1a3a2 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 21:09:22 -0800 Subject: [PATCH 09/12] add AbortReq.matches; rename PD_PREFILLING to PREFILL_LAUNCHED; make release_aborted_request public; read checker lists by kind; fix moved abort comments --- .../sglang/srt/constrained/grammar_manager.py | 4 +- python/sglang/srt/disaggregation/decode.py | 8 ++-- .../srt/disaggregation/encoder/receiver.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 10 ++--- python/sglang/srt/dllm/mixin/scheduler.py | 6 +-- python/sglang/srt/managers/io_struct.py | 4 ++ python/sglang/srt/managers/scheduler.py | 38 +++++++++---------- .../deferred_responses.py | 2 +- .../scheduler_components/invariant_checker.py | 9 ++--- .../scheduler_components/request_list.py | 2 +- python/sglang/test/scheduler_test_utils.py | 2 +- .../test_disaggregation_optimistic_prefill.py | 4 +- .../test_admission_abort_not_enqueued.py | 4 +- .../test_prefill_abort_result_cleanup.py | 2 +- .../test_prefill_complete_integration.py | 12 +++--- .../test_prefill_result_polling.py | 6 +-- .../unit/managers/test_auxiliary_output.py | 2 +- .../unit/managers/test_deferred_output.py | 2 +- .../test_disagg_idle_step_counters.py | 2 +- .../test_scheduler_chunked_abort_race.py | 4 +- ...est_scheduler_flush_cache_after_retract.py | 2 +- .../managers/test_scheduler_on_idle_load.py | 2 +- .../test_scheduler_pause_generation.py | 2 +- .../managers/test_scheduler_running_reqs.py | 2 - .../unit/managers/test_scheduler_timeouts.py | 4 +- .../test_scheduler_weight_version_tracking.py | 2 +- 26 files changed, 67 insertions(+), 72 deletions(-) diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index cbadb18d9bfa..990d52ecec13 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -116,11 +116,11 @@ def _pp_sync_ready_failed( return data def _abort(self, recv_req: AbortReq) -> None: - # Abort method 2: call `set_finish_with_abort` + # Call `set_finish_with_abort`. # The request will still run one prefill forward pass. # In this case, we change the input_ids to be only one token to make this prefill cheap. 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() diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d3f40669df65..6de87e7c8029 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1068,9 +1068,7 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: 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.abort_all or req.rid.startswith(recv_req.rid) - ): + for req in self.retracted_queue.take(lambda req: recv_req.matches(req.rid)): discard_kv_cache_backup( req, self.tree_cache, @@ -1081,10 +1079,10 @@ def _abort_retracted(self, recv_req: AbortReq) -> None: 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.abort_all or req.rid.startswith(recv_req.rid) + lambda req: recv_req.matches(req.rid) ): req.finish_outside_step(FINISH_ABORT()) - self.scheduler._release_aborted_request(req) + self.scheduler.release_aborted_request(req) self.scheduler.send_abort_output( make_abort_req(req, finished_reason=recv_req.finished_reason), req ) diff --git a/python/sglang/srt/disaggregation/encoder/receiver.py b/python/sglang/srt/disaggregation/encoder/receiver.py index a3818f765bfd..4c8305b3c0ff 100644 --- a/python/sglang/srt/disaggregation/encoder/receiver.py +++ b/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, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1132f0fece31..8a5b351a7b3f 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -484,9 +484,9 @@ def extend(self, reqs: List[Req], num_kv_heads: int) -> None: def _abort(self, recv_req: AbortReq) -> None: # Abort requests that have not yet been bootstrapped for req in self.queue: - if recv_req.abort_all or req.rid.startswith(recv_req.rid): + if recv_req.matches(req.rid): logger.debug(f"Abort bootstrap queue request. {req.rid=}") - self.scheduler._release_aborted_request(req) + self.scheduler.release_aborted_request(req) if hasattr(req.disagg_kv_sender, "abort"): req.disagg_kv_sender.abort() @@ -1219,9 +1219,9 @@ def advance_logprob_pt(i: int, req: Req) -> None: self.maybe_send_health_check_signal() def _abort_prefill_transfer(self: Scheduler, recv_req: AbortReq) -> None: - # Abort in-flight requests + # Abort requests whose KV is being sent for req in self.disagg_prefill_transfer_queue: - if recv_req.abort_all or req.rid.startswith(recv_req.rid): + if recv_req.matches(req.rid): logger.debug(f"Abort transfer queue request. {req.rid=}") if hasattr(req.disagg_kv_sender, "abort"): req.disagg_kv_sender.abort() @@ -1397,7 +1397,7 @@ def handle_inflight_transfer_failure( req.finish_outside_step(FINISH_LENGTH(length=0)) # unlock the tree release_kv_cache(req, self.tree_cache, checkpoint=True) - self._release_aborted_request(req) + self.release_aborted_request(req) if not isinstance(req.finished_reason, FINISH_ABORT): prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 8647c5576e4d..75bc308133cd 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -47,7 +47,7 @@ def _release_aborted_dllm_req(self: Scheduler, req: Req) -> None: # Finished in place, so a later abort in the same intake skips it. if req.finished_reason is None: req.finish_outside_step(FINISH_ABORT()) - self._release_aborted_request(req) + self.release_aborted_request(req) self.send_abort_output(make_abort_req(req), req) if req.kv.holds_kv or req.kv.holds_mamba: release_kv_cache(req, self.tree_cache, checkpoint=False) @@ -508,9 +508,7 @@ def filter_finished_reqs(self) -> None: self.staging_queue = [req for req in self.staging_queue if not req.finished()] def _abort(self, recv_req: AbortReq) -> None: - aborted_reqs = self.waiting_queue.take( - lambda req: recv_req.abort_all or req.rid.startswith(recv_req.rid) - ) + aborted_reqs = self.waiting_queue.take(lambda req: recv_req.matches(req.rid)) aborted = set(aborted_reqs) self.staging_queue = [req for req in self.staging_queue if req not in aborted] for req in aborted_reqs: diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 592d97fdceb4..be5db59ebce3 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2366,6 +2366,10 @@ def __post_init__(self): if self.rid is None: self.rid = "" + def matches(self, rid: str) -> bool: + """Whether this abort targets the request `rid`.""" + return self.abort_all or rid.startswith(self.rid) + class EncoderDispatchErrorReq(BaseReq, kw_only=True): """Tokenizer-to-scheduler failure for one EPD encoder dispatch.""" diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4a3890e3f3fd..4150545976be 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1388,7 +1388,7 @@ def init_running_status(self): # Prefill requests out of the chunked slot whose launched steps their next # holder has not taken yet: the running batch at the merge, the PD # transfer queue, HiSparse staging, or a requeue. - self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + self.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) # The current forward batch self.cur_batch_for_debug: Optional[ScheduleBatch] = None # The last forward batch @@ -2664,8 +2664,6 @@ def init_invariant_checker(self) -> None: get_last_batch=lambda: self.last_batch, get_request_lists=self.request_lists, get_chunked_req=lambda: self.chunked_req, - get_running_reqs=lambda: self.running_reqs, - get_prefill_launched_reqs=lambda: self.prefill_launched_reqs, get_running_batches=lambda: ( self.running_batch, *(self.running_mbs if get_parallel().pp_size > 1 else ()), @@ -3643,7 +3641,7 @@ def retire_unadmitted_request(self, req: Req) -> None: # `prefetch_from_storage` arms the paced-retry set for this attempt's # cache handle. Only a `finish`/ABORT and a `waiting_queue` sweep clear # that, and a retired request reaches neither. - self._release_aborted_request(req) + self.release_aborted_request(req) # `update_finish_state` returns early once `finished()`, so an already # set `finished_reason` is what the client receives; report the same. reason = req.finished_reason or req.to_finish @@ -3653,7 +3651,7 @@ def retire_unadmitted_request(self, req: Req) -> None: def _add_request_to_queue(self, req: Req, is_retracted: bool = False): if not self._set_or_validate_priority(req): - self._release_aborted_request(req) + self.release_aborted_request(req) return if is_retracted: req.storage_prefetch_retry_attempts = 0 @@ -3661,7 +3659,7 @@ def _add_request_to_queue(self, req: Req, is_retracted: bool = False): req.staged_prefetch_plan = None if self.disaggregation_mode == DisaggregationMode.NULL: if self._abort_on_queued_limit(req): - self._release_aborted_request(req) + self.release_aborted_request(req) return self._prefetch_kvcache(req) if self.uocr_runtime is None: @@ -3721,7 +3719,7 @@ def _set_or_validate_priority(self, req: Req) -> bool: return False return True - def _release_aborted_request(self, req: Req) -> None: + def release_aborted_request(self, req: Req) -> None: """Drop the cache-side state an aborted request left behind.""" req.release_mm_inputs_on_abort() self.tree_cache.finish(req.cache_request_handle, CacheRequestOutcome.ABORT) @@ -3738,7 +3736,7 @@ def _abort_uocr_doc_windows( ) -> None: """Abort document windows held off the waiting queue.""" for req in reqs: - self._release_aborted_request(req) + self.release_aborted_request(req) self.ipc_channels.send_to_tokenizer.send_output( make_abort_req(req, finished_reason=finished_reason), req ) @@ -3769,7 +3767,7 @@ def _abort_on_queued_limit(self, recv_req: Req) -> bool: direction * recv_req.priority < direction * candidate_req.priority ) if abort_existing_req: - self._release_aborted_request(candidate_req) + self.release_aborted_request(candidate_req) self.waiting_queue.pop(idx) self.beam_coordinator.retire_group(candidate_req) req_to_abort = candidate_req @@ -4014,7 +4012,7 @@ def process_pending_chunked_abort(self) -> None: req, self.req_to_metadata_buffer_idx_allocator ) req.pending_bootstrap = False - self._release_aborted_request(req) + self.release_aborted_request(req) # No member rows exist before the leader's prefill ends. self.beam_coordinator.retire_group(req) release_kv_cache(req, self.tree_cache, checkpoint=False) @@ -4689,7 +4687,7 @@ def request_lists(self) -> Dict[ListKind, Iterable[Req]]: ListKind.WAITING: self.waiting_queue, ListKind.GRAMMAR: self.grammar_manager.grammar_queue, ListKind.CHUNKED: [] if self.chunked_req is None else [self.chunked_req], - ListKind.PD_PREFILLING: self.prefill_launched_reqs, + ListKind.PREFILL_LAUNCHED: self.prefill_launched_reqs, ListKind.DECODING: self.running_reqs, } # Read at call time: a PD role switch rebuilds the PD queues. @@ -5691,7 +5689,7 @@ def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool: busy_lists = lists.keys() - { # Their steps are counted through the batches above. ListKind.DECODING, - ListKind.PD_PREFILLING, + ListKind.PREFILL_LAUNCHED, # Only its staging part is counted, above. ListKind.DLLM, # Held across a weight update, whose cache flush needs idle. @@ -6021,7 +6019,7 @@ def _collect_inflight_batches(self) -> list[ScheduleBatch | None]: def abort_request(self, recv_req: AbortReq): self.deferred_responses.abort(recv_req) if (chunked_req := self.chunked_req) is not None: - if recv_req.abort_all or chunked_req.rid.startswith(recv_req.rid): + if recv_req.matches(chunked_req.rid): self._pending_chunked_abort_req = chunked_req # todo hisparse, release resources for abort requests in hisparse coordinator @@ -6032,7 +6030,7 @@ def abort_request(self, recv_req: AbortReq): # Delete requests in the waiting queue to_del = [] for i, req in enumerate(self.waiting_queue): - if recv_req.abort_all or req.rid.startswith(recv_req.rid): + if recv_req.matches(req.rid): to_del.append(i) if self.uocr_runtime is not None: @@ -6053,7 +6051,7 @@ def abort_request(self, recv_req: AbortReq): # Finished in place, so a later abort in the same intake skips it. if req.finished_reason is None: req.finish_outside_step(FINISH_ABORT()) - self._release_aborted_request(req) + self.release_aborted_request(req) self.beam_coordinator.retire_group(req) # Without the initiator's reason the tokenizer falls back to a # generic abort message. @@ -6086,7 +6084,7 @@ def abort_request(self, recv_req: AbortReq): if self.disaggregation_mode == DisaggregationMode.DECODE: # Abort requests that have not yet finished preallocation for decode_req in self.disagg_decode_prealloc_queue.queue: - if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): + if recv_req.matches(decode_req.req.rid): logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") decode_req.kv_receiver.abort() if get_parallel().pp_size > 1: @@ -6094,7 +6092,7 @@ def abort_request(self, recv_req: AbortReq): # Abort requests waiting for kvcache to release tree cache for decode_req in self.disagg_decode_transfer_queue.queue: - if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): + if recv_req.matches(decode_req.req.rid): logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") if decode_req.host_staged: # Keep the host destination alive until prefill stops writing. @@ -6107,10 +6105,8 @@ def abort_request(self, recv_req: AbortReq): # Delete requests in the running batch for req in self.collect_inflight_reqs(): - if not req.finished() and ( - recv_req.abort_all or req.rid.startswith(recv_req.rid) - ): - # Abort method 3: `request_finish` + if not req.finished() and recv_req.matches(req.rid): + # Abort method 2: `request_finish` # The request will still run one decode forward pass. # Then we reuse all existing code to clean up the KV cache allocation. logger.debug(f"Abort running request. {req.rid=}") diff --git a/python/sglang/srt/managers/scheduler_components/deferred_responses.py b/python/sglang/srt/managers/scheduler_components/deferred_responses.py index 21671819cde5..2dbc2aff590b 100644 --- a/python/sglang/srt/managers/scheduler_components/deferred_responses.py +++ b/python/sglang/srt/managers/scheduler_components/deferred_responses.py @@ -86,7 +86,7 @@ def _abort_response( def abort(self, recv_req: AbortReq) -> None: for rid, (req, completion) in list(self.held.items()): - if recv_req.abort_all or rid.startswith(recv_req.rid): + if recv_req.matches(rid): del self.held[rid] try: completion.cancel() diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 0b16c0bdfceb..7abb28053b11 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -76,8 +76,6 @@ class SchedulerInvariantChecker: default=lambda: {} ) get_chunked_req: Callable = field(default=lambda: None) - get_running_reqs: Callable[[], Iterable[Req]] = field(default=lambda: ()) - get_prefill_launched_reqs: Callable[[], Iterable[Req]] = field(default=lambda: ()) get_running_batches: Callable = field(default=lambda: ()) # Under PP, the micro-batches' launched steps. get_pp_mbs: Callable = field(default=lambda: ()) @@ -487,8 +485,9 @@ def maybe_check_request_lists(self, batch: ScheduleBatch) -> None: live request is in exactly one list, one that matches its state.""" if not envs.SGLANG_ENABLE_RUNNING_REQS_CHECK.get(): return + lists = self.get_request_lists() owners: Dict[Req, ListKind] = {} - for kind, reqs in self.get_request_lists().items(): + for kind, reqs in lists.items(): for req in reqs: if req in owners: raise RuntimeError( @@ -509,14 +508,14 @@ def live(reqs: Iterable[Req]) -> Set[Req]: in_step = { req for b in (batch, *self.get_pp_mbs()) if b is not None for req in b.reqs } - running = set(self.get_running_reqs()) + running = set(lists.get(ListKind.DECODING, ())) _raise_if_any( live(in_running_batch | in_step) - owners.keys(), "Requests held by a running batch or a step but in no list", ) _raise_if_any( - in_running_batch & set(self.get_prefill_launched_reqs()), + in_running_batch & set(lists.get(ListKind.PREFILL_LAUNCHED, ())), "Requests in prefill_launched but in a running batch", ) _raise_if_any( diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index e275e35477f0..1aca947d80f8 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -27,7 +27,7 @@ class ListKind(enum.Enum): CHUNKED = enum.auto() PD_BOOTSTRAP = enum.auto() # Prefill requests out of the chunked slot whose launched steps are not done. - PD_PREFILLING = enum.auto() + PREFILL_LAUNCHED = enum.auto() PD_SENDING = enum.auto() PD_PREALLOC = enum.auto() PD_RETRACTED = enum.auto() diff --git a/python/sglang/test/scheduler_test_utils.py b/python/sglang/test/scheduler_test_utils.py index 2d777c845dbb..74095c260ff0 100644 --- a/python/sglang/test/scheduler_test_utils.py +++ b/python/sglang/test/scheduler_test_utils.py @@ -34,7 +34,7 @@ def empty_request_lists(scheduler: Scheduler) -> None: grammar_queue=_RequestList(ListKind.GRAMMAR) ) scheduler.chunked_req = None - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.running_reqs = _RequestList(ListKind.DECODING) scheduler.disagg_prefill_bootstrap_queue = None scheduler.disagg_prefill_transfer_queue = _RequestList(ListKind.PD_SENDING) diff --git a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py index 64a4b21cd6b9..c6bf451d5aa0 100644 --- a/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py +++ b/test/registered/disaggregation/test_disaggregation_optimistic_prefill.py @@ -477,9 +477,9 @@ def test_retry_release_retains_donated_checkpoint(self): disagg_prefill_bootstrap_queue=SimpleNamespace(queue=[]), metrics_reporter=SimpleNamespace(enable_metrics=False), processed_tokens_counter=0, - _release_aborted_request=lambda req: None, + release_aborted_request=lambda req: None, clear_pending_chunk_send=lambda req: None, - prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), + prefill_launched_reqs=_RequestList(ListKind.PREFILL_LAUNCHED), ) with patch( "sglang.srt.disaggregation.prefill.get_disagg", diff --git a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py index 5883bc5c2580..36c378ce0c46 100644 --- a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py +++ b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py @@ -71,7 +71,7 @@ def _make_scheduler(): sched = SimpleNamespace( output_streamer=MagicMock(), beam_coordinator=MagicMock(), - _release_aborted_request=MagicMock(), + release_aborted_request=MagicMock(), ) sched.retire_unadmitted_request = MagicMock( side_effect=lambda req: Scheduler.retire_unadmitted_request(sched, req) @@ -315,7 +315,7 @@ def test_retire_detaches_a_streaming_session(self): # PREFILL arms the cache's paced-retry set via `_prefetch_kvcache` # before the door, and only this call clears it again. - sched._release_aborted_request.assert_called_once_with(req) + sched.release_aborted_request.assert_called_once_with(req) session.abort_req.assert_called_once() self.assertIsNone(req.session) sched.beam_coordinator.retire_group.assert_called_once_with(req) diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index bb3f4c8bc2b0..195694d0d977 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -60,7 +60,7 @@ def update_finish_state(self): class _Scheduler(SchedulerDisaggregationPrefillMixin): def __init__(self): - self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + self.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) self.batch_result_processor = SimpleNamespace( snapshot_auxiliary_output_starts=Mock(return_value=[]), move_logprobs_to_cpu=Mock(), diff --git a/test/registered/unit/disaggregation/test_prefill_complete_integration.py b/test/registered/unit/disaggregation/test_prefill_complete_integration.py index db0a1b607181..0865e8159453 100644 --- a/test/registered/unit/disaggregation/test_prefill_complete_integration.py +++ b/test/registered/unit/disaggregation/test_prefill_complete_integration.py @@ -352,7 +352,7 @@ def test_exhausted_optimistic_attempt_follows_allocation_policy(self): scheduler = SimpleNamespace( tree_cache=Mock(supports_mamba=Mock(return_value=False)), clear_pending_chunk_send=Mock(), - prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), + prefill_launched_reqs=_RequestList(ListKind.PREFILL_LAUNCHED), waiting_queue=[], disagg_prefill_bootstrap_queue=SimpleNamespace( queue=_RequestList(ListKind.PD_BOOTSTRAP) @@ -438,8 +438,8 @@ class _Req: metrics_reporter=SimpleNamespace(enable_metrics=False), processed_tokens_counter=0, ) - scheduler._release_aborted_request = MethodType( - Scheduler._release_aborted_request, scheduler + scheduler.release_aborted_request = MethodType( + Scheduler.release_aborted_request, scheduler ) with ( patch("sglang.srt.disaggregation.prefill.checkpoint_kv_cache"), @@ -476,7 +476,9 @@ def test_chunked_prefill_does_not_yield_for_decode_admission(self): check_bootstrap=Mock(return_value=False), has_bootstrapped_waiting_req=Mock(return_value=True), optimistic_release_and_requeue=Mock(), - prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), + prefill_launched_reqs=_RequestList( + ListKind.PREFILL_LAUNCHED + ), ) scheduler.checkpoint_disagg_prefill = MethodType( SchedulerDisaggregationPrefillMixin.checkpoint_disagg_prefill, @@ -533,7 +535,7 @@ def test_yielded_request_requeues_after_its_chunk_once_bootstrap_completes(self) ), spec_algorithm=SimpleNamespace(is_eagle=lambda: False), optimistic_release_and_requeue=Mock(), - prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), + prefill_launched_reqs=_RequestList(ListKind.PREFILL_LAUNCHED), metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()), maybe_send_health_check_signal=Mock(), ) diff --git a/test/registered/unit/disaggregation/test_prefill_result_polling.py b/test/registered/unit/disaggregation/test_prefill_result_polling.py index 91b1de000f77..a08ecce39720 100644 --- a/test/registered/unit/disaggregation/test_prefill_result_polling.py +++ b/test/registered/unit/disaggregation/test_prefill_result_polling.py @@ -185,7 +185,7 @@ def __init__( self.decode_offload_manager = None self.running_batch = ScheduleBatch(reqs=[]) self.running_reqs = _RequestList(ListKind.DECODING) - self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + self.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) self.invariant_checker = Mock() self.uocr_runtime = None self.waiting_queue = [] @@ -1160,7 +1160,7 @@ def test_retract_resume_requeues_optimistic_chunk_after_result_finishes(self): ) class OptimisticScheduler(_Scheduler): - _release_aborted_request = Scheduler._release_aborted_request + release_aborted_request = Scheduler.release_aborted_request def run_batch(self, batch): if req in batch.reqs: @@ -1275,7 +1275,7 @@ def __init__(self): self.disaggregation_mode = DisaggregationMode.PREFILL self.dllm_config = None self.running_reqs = _RequestList(ListKind.DECODING) - self.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + self.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) self.enable_staging = False self.enable_overlap = False self.token_to_kv_pool_allocator = SimpleNamespace(page_size=64) diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 77e05458ef0e..00c546de9ef5 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -554,7 +554,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): spec_algorithm=SimpleNamespace(is_eagle=lambda: False), tree_cache=object(), disagg_prefill_transfer_queue=_RequestList(ListKind.PD_SENDING), - prefill_launched_reqs=_RequestList(ListKind.PD_PREFILLING), + prefill_launched_reqs=_RequestList(ListKind.PREFILL_LAUNCHED), checkpoint_disagg_prefill=Mock(), send_kv_chunk=Mock(), metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()), diff --git a/test/registered/unit/managers/test_deferred_output.py b/test/registered/unit/managers/test_deferred_output.py index 30e3f33d52e4..15ab32dd27a0 100644 --- a/test/registered/unit/managers/test_deferred_output.py +++ b/test/registered/unit/managers/test_deferred_output.py @@ -57,7 +57,7 @@ def _scheduler(): scheduler.chunked_req = None scheduler.hisparse_coordinator = None scheduler.collect_inflight_reqs = lambda: set() - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_disagg_idle_step_counters.py b/test/registered/unit/managers/test_disagg_idle_step_counters.py index 69f852fad413..f9f56c75dcc5 100644 --- a/test/registered/unit/managers/test_disagg_idle_step_counters.py +++ b/test/registered/unit/managers/test_disagg_idle_step_counters.py @@ -547,7 +547,7 @@ def record_result(self, scheduler, batch, result_ts, *, result=None): def make_scheduler(self, schedule): scheduler = Scheduler.__new__(Scheduler) scheduler.running_reqs = _RequestList(ListKind.DECODING) - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.tp_size = get_parallel().tp_size scheduler.gracefully_exit = False scheduler._engine_paused = False diff --git a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py index f528f9c4f5c8..1d148c010a0e 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_abort_race.py +++ b/test/registered/unit/managers/test_scheduler_chunked_abort_race.py @@ -81,7 +81,7 @@ def test_discarded_chunk_keeps_its_abort_reason(self): req = _FakeReq("failed_load") req.discard_output_reason = FINISH_ABORT("External KV cache load failed", 503) sched = _make_scheduler(req, chunked_req=req, running_reqs=[]) - sched._release_aborted_request = Mock() + sched.release_aborted_request = Mock() sched.tree_cache = None sched.ipc_channels = Mock() @@ -95,7 +95,7 @@ def test_discarded_chunk_keeps_its_abort_reason(self): def test_aborted_chunked_beam_leader_retires_its_group(self): req = _FakeReq("beam_leader") sched = _make_scheduler(req, chunked_req=req, running_reqs=[]) - sched._release_aborted_request = Mock() + sched.release_aborted_request = Mock() sched.tree_cache = None sched.ipc_channels = Mock() diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py index 8092ef19578b..45d8fbaa4129 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py +++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py @@ -63,7 +63,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.enable_hierarchical_cache = False scheduler.enable_lmcache = False scheduler.enable_hisparse = False - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_scheduler_on_idle_load.py b/test/registered/unit/managers/test_scheduler_on_idle_load.py index d0f3f562848a..12ab08d0056d 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/test/registered/unit/managers/test_scheduler_on_idle_load.py @@ -123,7 +123,7 @@ def _idle_scheduler(self, materializer): scheduler.grammar_manager = SimpleNamespace(grammar_queue=[]) scheduler.disaggregation_mode = DisaggregationMode.NULL scheduler.enable_hisparse = False - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.running_reqs = [] scheduler.dllm_config = None scheduler.disagg_prefill_bootstrap_queue = None diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 547bb8a27068..8f0e81aed081 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -102,7 +102,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.cur_batch_for_debug = None scheduler.chunked_req = None scheduler.running_reqs = _RequestList(ListKind.DECODING) - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.running_batch = MagicMock() scheduler.running_batch.reqs = [] scheduler.running_batch.is_empty.return_value = True diff --git a/test/registered/unit/managers/test_scheduler_running_reqs.py b/test/registered/unit/managers/test_scheduler_running_reqs.py index 1ac9a3780638..9b876abe3b7a 100644 --- a/test/registered/unit/managers/test_scheduler_running_reqs.py +++ b/test/registered/unit/managers/test_scheduler_running_reqs.py @@ -86,8 +86,6 @@ def _invariant_checker(s): scheduler_stage_metrics=None, get_request_lists=s.request_lists, get_chunked_req=lambda: s.chunked_req, - get_running_reqs=lambda: s.running_reqs, - get_prefill_launched_reqs=lambda: s.prefill_launched_reqs, get_running_batches=lambda: (s.running_batch,), ) diff --git a/test/registered/unit/managers/test_scheduler_timeouts.py b/test/registered/unit/managers/test_scheduler_timeouts.py index 37058db38891..96e10fc41721 100644 --- a/test/registered/unit/managers/test_scheduler_timeouts.py +++ b/test/registered/unit/managers/test_scheduler_timeouts.py @@ -168,14 +168,14 @@ def test_popped_request_is_finished_and_a_later_abort_skips_it(self): s.mm_receiver = None s.dllm_config = None s.grammar_manager = MagicMock() - s._release_aborted_request = MagicMock() + s.release_aborted_request = MagicMock() s.abort_request(AbortReq(rid="queued")) s.abort_request(AbortReq(abort_all=True)) self.assertIsInstance(req.finished_reason, FINISH_ABORT) self.assertIsNone(req.to_finish) - s._release_aborted_request.assert_called_once_with(req) + s.release_aborted_request.assert_called_once_with(req) class TestRunningTimeout(CustomTestCase): diff --git a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py index 3db4629fe735..0d5ca0720d66 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -173,7 +173,7 @@ def _scheduler(self) -> Scheduler: scheduler.load_inquirer = MagicMock() scheduler.tree_cache = SimpleNamespace(flush_pending_backups=MagicMock()) scheduler.running_reqs = _RequestList(ListKind.DECODING) - scheduler.prefill_launched_reqs = _RequestList(ListKind.PD_PREFILLING) + scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED) scheduler.invariant_checker = MagicMock() scheduler.kv_weight_version_tracker = KvWeightVersionTracker( num_slots=8, From d353f9eaebf3c7bd32b7648bb87cea4af487ac1b Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 21:11:53 -0800 Subject: [PATCH 10/12] trim moved abort comments --- python/sglang/srt/constrained/grammar_manager.py | 5 ++--- python/sglang/srt/disaggregation/prefill.py | 2 -- .../sglang/srt/managers/scheduler_components/request_list.py | 1 - 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 990d52ecec13..e3f82fdb5d7e 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -116,9 +116,8 @@ def _pp_sync_ready_failed( return data def _abort(self, recv_req: AbortReq) -> None: - # Call `set_finish_with_abort`. - # The request will still run one prefill forward pass. - # In this case, we change the input_ids to be only one token to make this prefill cheap. + # 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.matches(req.rid): logger.debug(f"Abort grammar queue request. {req.rid=}") diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8a5b351a7b3f..3a3f47fae3ac 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -482,7 +482,6 @@ def extend(self, reqs: List[Req], num_kv_heads: int) -> None: self.add(req, num_kv_heads) def _abort(self, recv_req: AbortReq) -> None: - # Abort requests that have not yet been bootstrapped for req in self.queue: if recv_req.matches(req.rid): logger.debug(f"Abort bootstrap queue request. {req.rid=}") @@ -1219,7 +1218,6 @@ def advance_logprob_pt(i: int, req: Req) -> None: self.maybe_send_health_check_signal() def _abort_prefill_transfer(self: Scheduler, recv_req: AbortReq) -> None: - # Abort requests whose KV is being sent for req in self.disagg_prefill_transfer_queue: if recv_req.matches(req.rid): logger.debug(f"Abort transfer queue request. {req.rid=}") diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index 1aca947d80f8..d92f4492b748 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -26,7 +26,6 @@ class ListKind(enum.Enum): DECODING = enum.auto() CHUNKED = enum.auto() PD_BOOTSTRAP = enum.auto() - # Prefill requests out of the chunked slot whose launched steps are not done. PREFILL_LAUNCHED = enum.auto() PD_SENDING = enum.auto() PD_PREALLOC = enum.auto() From 4d60fc3434e4a8e5ae68779d3c7cc1f2abbfc0d0 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 21:14:23 -0800 Subject: [PATCH 11/12] build real AbortReq in grammar abort tests --- .../unit/constrained/test_grammar_manager.py | 21 ++++++------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index c28179394d35..c3c3961a6da4 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -35,6 +35,7 @@ ) from sglang.srt.constrained.xgrammar_backend import XGrammarGrammarBackend from sglang.srt.distributed.communication_tags import P2PTag +from sglang.srt.managers.io_struct import AbortReq from sglang.srt.parser.inkling_tokenizer import ( CONTENT_MODEL_END_SAMPLING, CONTENT_THINKING, @@ -487,9 +488,7 @@ def test_abort_by_rid_prefix(self): req.grammar = future mgr.grammar_queue.join([req]) - abort_req = MagicMock() - abort_req.abort_all = False - abort_req.rid = "req-123" + abort_req = AbortReq(rid="req-123", abort_all=False) mgr.grammar_queue.abort(abort_req) future.cancel.assert_called_once() @@ -501,9 +500,7 @@ def test_abort_non_matching_rid(self): req.grammar = MagicMock(spec=Future) mgr.grammar_queue.join([req]) - abort_req = MagicMock() - abort_req.abort_all = False - abort_req.rid = "req-123" + abort_req = AbortReq(rid="req-123", abort_all=False) mgr.grammar_queue.abort(abort_req) req.set_finish_with_abort.assert_not_called() @@ -517,9 +514,7 @@ def test_abort_all(self): mgr.grammar_queue.join([req]) reqs.append(req) - abort_req = MagicMock() - abort_req.abort_all = True - abort_req.rid = "" + abort_req = AbortReq(rid="", abort_all=True) mgr.grammar_queue.abort(abort_req) for req in reqs: @@ -528,9 +523,7 @@ def test_abort_all(self): def test_abort_empty_queue(self): """Aborting on an empty queue should not raise.""" mgr = self._make_mgr_with_queue() - abort_req = MagicMock() - abort_req.abort_all = True - abort_req.rid = "" + abort_req = AbortReq(rid="", abort_all=True) mgr.grammar_queue.abort(abort_req) # Should not raise def test_abort_prefix_match(self): @@ -540,9 +533,7 @@ def test_abort_prefix_match(self): req.grammar = MagicMock(spec=Future) mgr.grammar_queue.join([req]) - abort_req = MagicMock() - abort_req.abort_all = False - abort_req.rid = "req-123" + abort_req = AbortReq(rid="req-123", abort_all=False) mgr.grammar_queue.abort(abort_req) req.set_finish_with_abort.assert_called_once() From f3ec76a0eaa0e13a0f0c5ca198f65d6a0d7b3cb5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 22:30:46 -0800 Subject: [PATCH 12/12] use an identity-hashed fake request in the retraction backup test --- .../registered/unit/mem_cache/test_decode_retraction_backup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/registered/unit/mem_cache/test_decode_retraction_backup.py b/test/registered/unit/mem_cache/test_decode_retraction_backup.py index cf20d02a5550..65cee20f13c7 100644 --- a/test/registered/unit/mem_cache/test_decode_retraction_backup.py +++ b/test/registered/unit/mem_cache/test_decode_retraction_backup.py @@ -52,6 +52,7 @@ HiCacheDraftPlan, ) from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase from sglang.test.tree_slots import tree_slots_of @@ -247,7 +248,7 @@ def _build_cache( ) def _admit_req(self, env, num_tokens: int): - req = SimpleNamespace(rid="request", kv=ReqKvInfo(), seqlen=num_tokens + 1) + req = FakeReq(rid="request", kv=ReqKvInfo(), seqlen=num_tokens + 1) self.assertIsNotNone(env.req_to_token_pool.alloc([req])) source_indices = env.allocator.alloc(num_tokens) self.assertIsNotNone(source_indices)