[Fix] scheduler: Record weight version events on every request list

python/sglang/srt/managers/scheduler.py

@@ -2747,14 +2747,8 @@ def init_load_inquirer(self) -> None:
             ),
             get_stats=lambda: self.metrics_reporter.stats,
             get_chunked_req=lambda: self.chunked_req,
-            get_disagg_prefill_bootstrap_queue=lambda: (
-                self.disagg_prefill_bootstrap_queue
-            ),
-            get_disagg_prefill_transfer_queue=lambda: (
-                self.disagg_prefill_transfer_queue
-            ),
+            get_request_lists=self.request_lists,
             get_disagg_decode_prealloc_queue=lambda: self.disagg_decode_prealloc_queue,
-            get_disagg_decode_transfer_queue=lambda: self.disagg_decode_transfer_queue,
             get_spec_total_num_accept_tokens=lambda: (
                 self.metrics_reporter.spec_total_num_accept_tokens
             ),
@@ -5987,14 +5981,7 @@ def record_weight_version_change(self, new_version: Optional[str]) -> None:
         old_version = get_serving().weight_version
         get_context().override("scheduler.weight_version", weight_version=new_version)
 
-        live_reqs = {
-            *self.collect_inflight_reqs(),
-            *self.waiting_queue,
-            *([self.chunked_req] if self.chunked_req is not None else []),
-            *self.deferred_responses.held,
-        }
-        if self.hisparse_coordinator is not None:
-            live_reqs.update(self.hisparse_coordinator.ack_staging_queue)
+        live_reqs = [req for reqs in self.request_lists().values() for req in reqs]
         num_recorded = record_weight_version_events(live_reqs, old_version=old_version)
         logger.info(
             f"Weight version changed. {old_version=} {new_version=} {num_recorded=}"

python/sglang/srt/managers/scheduler_components/load_inquirer.py

@@ -14,6 +14,7 @@
     QueueMetrics,
     SpeculativeMetrics,
 )
+from sglang.srt.managers.scheduler_components.request_list import ListKind
 from sglang.srt.runtime_context import get_lora, get_parallel
 
 if TYPE_CHECKING:
@@ -46,10 +47,8 @@ class SchedulerLoadInquirer:
     get_recent_cache_hit_rate: Callable
     get_stats: Callable
     get_chunked_req: Callable
-    get_disagg_prefill_bootstrap_queue: Callable
-    get_disagg_prefill_transfer_queue: Callable
+    get_request_lists: Callable
     get_disagg_decode_prealloc_queue: Callable
-    get_disagg_decode_transfer_queue: Callable
     get_spec_total_num_accept_tokens: Callable
     get_spec_total_num_forward_ct: Callable
     get_total_prefill_uncached_tokens: Callable
@@ -111,19 +110,18 @@ def get_loads(self) -> LoadSnapshot:
         stats = self.get_stats()
         num_running_reqs = len(self.get_running_batch().reqs)
 
+        lists = self.get_request_lists()
         waiting_queues = [self.get_waiting_queue()]
         pending_token_queues = [self.get_waiting_queue()]
         awaiting_kv_tokens = 0
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            prefill_bootstrap_queue = self.get_disagg_prefill_bootstrap_queue().queue
+            prefill_bootstrap_queue = lists[ListKind.PD_BOOTSTRAP]
             waiting_queues.append(prefill_bootstrap_queue)
             pending_token_queues.append(prefill_bootstrap_queue)
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
-            decode_prealloc_queue = self.get_disagg_decode_prealloc_queue().queue
-            decode_transfer_queue = self.get_disagg_decode_transfer_queue().queue
-            decode_retracted_queue = (
-                self.get_disagg_decode_prealloc_queue().retracted_queue
-            )
+            decode_prealloc_queue = lists[ListKind.PD_PREALLOC]
+            decode_transfer_queue = lists[ListKind.PD_RECEIVING]
+            decode_retracted_queue = lists[ListKind.PD_RETRACTED]
             waiting_queues.append(decode_prealloc_queue)
             waiting_queues.append(decode_transfer_queue)
             waiting_queues.append(decode_retracted_queue)
@@ -188,15 +186,13 @@ def get_loads(self) -> LoadSnapshot:
         num_prealloc_ready_tokens = 0
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
             mode_str = "prefill"
-            prefill_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue)
-            prefill_inflight = len(self.get_disagg_prefill_transfer_queue())
+            prefill_bootstrap = len(lists[ListKind.PD_BOOTSTRAP])
+            prefill_inflight = len(lists[ListKind.PD_SENDING])
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             mode_str = "decode"
-            decode_prealloc = len(self.get_disagg_decode_prealloc_queue().queue)
-            decode_transfer = len(self.get_disagg_decode_transfer_queue().queue)
-            decode_retracted = len(
-                self.get_disagg_decode_prealloc_queue().retracted_queue
-            )
+            decode_prealloc = len(lists[ListKind.PD_PREALLOC])
+            decode_transfer = len(lists[ListKind.PD_RECEIVING])
+            decode_retracted = len(lists[ListKind.PD_RETRACTED])
             ready_reqs = [
                 decode_req.req
                 for decode_req in self.get_disagg_decode_prealloc_queue().queue.items()

python/sglang/srt/managers/scheduler_components/metrics_reporter.py

@@ -15,6 +15,7 @@
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import EPLB_BALANCEDNESS_WINDOW_SIZES
 from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler_components.request_list import ListKind
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.mem_cache.hicache_storage import PoolName
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
@@ -855,10 +856,7 @@ def report_prefill_stats(
                 )
             else:
                 self.stats.num_running_reqs = prefill_stats.num_running_reqs
-            self.stats.num_queue_reqs = QueueCount.from_reqs(
-                self.scheduler.waiting_queue, priority_enabled
-            )
-            self.stats.num_grammar_queue_reqs = len(self.scheduler.grammar_manager)
+            self._update_queue_stats(priority_enabled)
             self.stats.cache_hit_rate = cache_hit_rate
             # Refresh here too: prefill-heavy stretches can run long between
             # decode-stats ticks, and the gauge must decay rather than hold.
@@ -874,13 +872,6 @@ def report_prefill_stats(
 
             # PD disaggregation
             if self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
-                self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs(
-                    self.scheduler.disagg_prefill_bootstrap_queue.queue,
-                    priority_enabled,
-                )
-                self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
-                    self.scheduler.disagg_prefill_transfer_queue, priority_enabled
-                )
                 self.stats.kv_transfer_speed_gb_s = self.kv_transfer_speed_gb_s
                 self.stats.kv_transfer_latency_ms = self.kv_transfer_latency_ms
             elif self.scheduler.disaggregation_mode == DisaggregationMode.DECODE:
@@ -1073,10 +1064,7 @@ def report_decode_stats(
             self.stats.num_running_reqs = QueueCount.from_reqs(
                 batch.reqs, priority_enabled
             )
-            self.stats.num_queue_reqs = QueueCount.from_reqs(
-                self.scheduler.waiting_queue, priority_enabled
-            )
-            self.stats.num_grammar_queue_reqs = len(self.scheduler.grammar_manager)
+            self._update_queue_stats(priority_enabled)
             self.stats.gen_throughput = self.last_gen_throughput
             # cache_hit_rate is prefill-owned (per-report semantics); decode
             # ticks must not reset it, or the exported gauge reads 0 whenever
@@ -1099,16 +1087,6 @@ def report_decode_stats(
             self.stats.num_paused_reqs = self.num_paused_reqs
             self.num_retracted_reqs = self.num_paused_reqs = 0
 
-            # PD disaggregation
-            if self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
-                self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs(
-                    self.scheduler.disagg_prefill_bootstrap_queue.queue,
-                    priority_enabled,
-                )
-                self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
-                    self.scheduler.disagg_prefill_transfer_queue, priority_enabled
-                )
-
             # Streaming session metrics
             self.stats.num_streaming_sessions = (
                 self.scheduler.pool_stats_observer.streaming_session_count()
@@ -1437,6 +1415,20 @@ def _reset_device_timer_window(self):
             self.fwd_occupancy = float("nan")
             self.stats.fwd_occupancy = float("nan")
 
+    def _update_queue_stats(self, priority_enabled: bool) -> None:
+        lists = self.scheduler.request_lists()
+        self.stats.num_queue_reqs = QueueCount.from_reqs(
+            lists[ListKind.WAITING], priority_enabled
+        )
+        self.stats.num_grammar_queue_reqs = len(lists[ListKind.GRAMMAR])
+        if self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
+            self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs(
+                lists[ListKind.PD_BOOTSTRAP], priority_enabled
+            )
+            self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
+                lists[ListKind.PD_SENDING], priority_enabled
+            )
+
     def _update_decode_queue_stats(self, priority_enabled: bool) -> None:
         transfer_queue = self.scheduler.disagg_decode_transfer_queue.queue
         self.stats.num_decode_prealloc_queue_reqs = QueueCount.from_reqs(
@@ -1509,17 +1501,7 @@ def _maybe_log_idle_metrics(self):
             self.scheduler.running_batch.reqs, priority_enabled
         )
         self.stats.gen_throughput = 0
-        self.stats.num_queue_reqs = QueueCount.from_reqs(
-            self.scheduler.waiting_queue, priority_enabled
-        )
-        self.stats.num_grammar_queue_reqs = len(self.scheduler.grammar_manager)
-        if self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL:
-            self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs(
-                self.scheduler.disagg_prefill_bootstrap_queue.queue, priority_enabled
-            )
-            self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs(
-                self.scheduler.disagg_prefill_transfer_queue, priority_enabled
-            )
+        self._update_queue_stats(priority_enabled)
         if self.scheduler.disaggregation_mode == DisaggregationMode.DECODE:
             self._update_decode_queue_stats(priority_enabled)
         # Host pools change without a batch (e.g. /flush_cache); refresh before

test/registered/unit/managers/test_deferred_output.py

@@ -19,10 +19,6 @@
 from sglang.srt.managers.scheduler_components.output_streamer import (
     SchedulerOutputStreamer,
 )
-from sglang.srt.managers.scheduler_components.request_list import (
-    ListKind,
-    _RequestList,
-)
 from sglang.srt.managers.scheduler_components.request_receiver import (
     SchedulerRequestReceiver,
 )
@@ -33,6 +29,7 @@
     compute_weight_version_spans,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.scheduler_test_utils import empty_request_lists
 from sglang.test.test_utils import published_topology
 
 register_cpu_ci(est_time=6, suite="base-a-test-cpu")
@@ -54,17 +51,8 @@ def _req(rid, return_logprob=False):
 def _scheduler():
     scheduler = object.__new__(Scheduler)
     scheduler.disaggregation_mode = DisaggregationMode.NULL
-    scheduler.waiting_queue = []
-    scheduler.chunked_req = None
-    scheduler.hisparse_coordinator = None
+    empty_request_lists(scheduler)
     scheduler.collect_inflight_reqs = lambda: set()
-    scheduler.prefill_launched_reqs = _RequestList(ListKind.PREFILL_LAUNCHED)
-    scheduler.running_reqs = []
-    scheduler.dllm_config = None
-    scheduler.disagg_prefill_bootstrap_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()
     scheduler.rust_server = None
     streamer = object.__new__(SchedulerOutputStreamer)

test/registered/unit/managers/test_load_inquirer.py

@@ -70,10 +70,8 @@ def make_inquirer(
             get_chunked_req=lambda: chunked_req,
             get_recent_cache_hit_rate=lambda: 0.75,
             get_stats=None,
-            get_disagg_prefill_bootstrap_queue=None,
-            get_disagg_prefill_transfer_queue=None,
+            get_request_lists=None,
             get_disagg_decode_prealloc_queue=None,
-            get_disagg_decode_transfer_queue=None,
             get_spec_total_num_accept_tokens=None,
             get_spec_total_num_forward_ct=None,
             get_total_prefill_uncached_tokens=None,
@@ -136,12 +134,13 @@ def test_decode_loads_count_prealloc_requests_waiting_for_input(self):
             spec_algorithm=SimpleNamespace(is_none=lambda: True),
             get_running_batch=lambda: SimpleNamespace(reqs=[]),
             get_stats=SchedulerStats,
-            get_disagg_decode_prealloc_queue=lambda: SimpleNamespace(
-                queue=prealloc, retracted_queue=_RequestList(ListKind.PD_RETRACTED)
-            ),
-            get_disagg_decode_transfer_queue=lambda: SimpleNamespace(
-                queue=decode_queue(ListKind.PD_RECEIVING)
-            ),
+            get_disagg_decode_prealloc_queue=lambda: SimpleNamespace(queue=prealloc),
+            get_request_lists=lambda: {
+                ListKind.WAITING: [],
+                ListKind.PD_PREALLOC: prealloc,
+                ListKind.PD_RECEIVING: decode_queue(ListKind.PD_RECEIVING),
+                ListKind.PD_RETRACTED: _RequestList(ListKind.PD_RETRACTED),
+            },
             get_total_prefill_uncached_tokens=lambda: 0,
             get_total_prefill_busy_us=lambda: 0,
             get_decode_moment_totals=lambda: (0,),
@@ -238,6 +237,7 @@ def test_scheduler_wires_storage_provider(self):
             metrics_reporter=SimpleNamespace(recent_cache_hit_rate=0.75),
             waiting_queue=[req],
             chunked_req=None,
+            request_lists=None,
         )
 
         Scheduler.init_load_inquirer(scheduler)

test/registered/unit/managers/test_scheduler_weight_version_tracking.py

@@ -45,62 +45,6 @@ def __init__(self, weight_version: str):
         self.weight_version = weight_version
 
 
-class _ContextStub:
-    def __init__(self, serving: _ServingStub):
-        self.serving = serving
-
-    def override(self, source, **fields):
-        self.serving.weight_version = fields["weight_version"]
-
-
-class TestSchedulerRecordWeightVersionChange(CustomTestCase):
-    def _serving(self, version: str) -> _ServingStub:
-        serving = _ServingStub(version)
-        for name, value in (
-            ("get_serving", serving),
-            ("get_context", _ContextStub(serving)),
-        ):
-            patcher = patch(f"sglang.srt.managers.scheduler.{name}", return_value=value)
-            patcher.start()
-            self.addCleanup(patcher.stop)
-        return serving
-
-    def _scheduler(
-        self, *, inflight=(), waiting=(), chunked=None, staging=()
-    ) -> SimpleNamespace:
-        return SimpleNamespace(
-            collect_inflight_reqs=lambda: set(inflight),
-            deferred_responses=DeferredResponses(
-                disaggregation_mode=DisaggregationMode.NULL,
-                output_streamer=None,
-                send_abort=None,
-            ),
-            waiting_queue=list(waiting),
-            chunked_req=chunked,
-            hisparse_coordinator=(
-                SimpleNamespace(ack_staging_queue=list(staging)) if staging else None
-            ),
-        )
-
-    def test_every_source_of_live_requests_is_stamped(self):
-        """A request missed here keeps attributing its next tokens to the superseded version."""
-        self._serving("v1")
-        inflight, queued, chunked, staged = (object() for _ in range(4))
-        scheduler = self._scheduler(
-            inflight=[inflight], waiting=[queued], chunked=chunked, staging=[staged]
-        )
-
-        with patch(
-            "sglang.srt.managers.scheduler.record_weight_version_events",
-            return_value=0,
-        ) as recorder:
-            Scheduler.record_weight_version_change(scheduler, new_version="v2")
-
-        self.assertEqual(
-            set(recorder.call_args.args[0]), {inflight, queued, chunked, staged}
-        )
-
-
 class TestSchedulerBatchWeightVersion(CustomTestCase):
     def test_embedding_rejects_enabled_prefill_tracking(self) -> None:
         """Embedding and reward models reject enabled prefill tracking before allocation."""

test/registered/unit/observability/test_forward_pass_metrics.py

@@ -1021,8 +1021,7 @@ class TestIdleMetrics(CustomTestCase):
     def setUp(self):
         self.scheduler = types.SimpleNamespace(
             running_batch=types.SimpleNamespace(reqs=[]),
-            waiting_queue=[],
-            grammar_manager=[],
+            request_lists=lambda: {kind: [] for kind in ListKind},
             enable_priority_scheduling=False,
             disaggregation_mode=DisaggregationMode.NULL,
             pool_stats_observer=types.SimpleNamespace(
@@ -1046,8 +1045,6 @@ def setUp(self):
 
     def test_idle_publishes_refreshed_kv_metrics(self):
         self.reporter.current_scheduler_metrics_enabled = True
-        self.scheduler.disagg_prefill_bootstrap_queue = types.SimpleNamespace(queue=[])
-        self.scheduler.disagg_prefill_transfer_queue = []
         self.scheduler.disagg_decode_prealloc_queue = types.SimpleNamespace(
             queue=[], pre_allocated_token_counts=lambda: (0, 0)
         )

test/registered/unit/observability/test_stat_loggers_di.py

@@ -30,6 +30,7 @@
 from sglang.srt.managers.scheduler_components.metrics_reporter import (
     SchedulerMetricsReporter,
 )
+from sglang.srt.managers.scheduler_components.request_list import ListKind
 from sglang.srt.mem_cache.hicache_storage import PoolName
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.observability.metrics_collector import (
@@ -283,8 +284,7 @@ def test_logs_current_deferred_release_count(self):
                 queue=[], pre_allocated_token_counts=lambda: (0, 0)
             ),
             running_batch=SimpleNamespace(reqs=[]),
-            waiting_queue=[],
-            grammar_manager=[],
+            request_lists=lambda: {kind: [] for kind in ListKind},
             enable_fpm=False,
             enable_priority_scheduling=False,
             kv_events_publisher=SimpleNamespace(emit_kv_metrics=MagicMock()),
@@ -367,7 +367,7 @@ def test_decode_report_logs_current_deferred_release_count(self):
             ),
             running_batch=batch,
             waiting_queue=[],
-            grammar_manager=[],
+            request_lists=lambda: {kind: [] for kind in ListKind},
             enable_priority_scheduling=False,
             max_total_num_tokens=1,
             spec_algorithm=SimpleNamespace(is_none=lambda: True),
@@ -576,8 +576,7 @@ def test_idle_report_refreshes_the_host_pools(self):
         )
         reporter.scheduler.disaggregation_mode = DisaggregationMode.NULL
         reporter.scheduler.running_batch = SimpleNamespace(reqs=[])
-        reporter.scheduler.waiting_queue = []
-        reporter.scheduler.grammar_manager = []
+        reporter.scheduler.request_lists = lambda: {kind: [] for kind in ListKind}
         reporter.scheduler.enable_fpm = False
         reporter.scheduler.enable_priority_scheduling = False
         reporter.scheduler.pool_stats_observer = SimpleNamespace(