From a7af2392591c121ce7b287baa0ce6e8deaa02bc5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 21:53:31 -0800 Subject: [PATCH 1/5] record weight version events on every request list --- python/sglang/srt/managers/scheduler.py | 9 +- .../test_scheduler_weight_version_tracking.py | 31 ++---- .../unit/utils/test_weight_versions.py | 98 +++++-------------- 3 files changed, 36 insertions(+), 102 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 171199a00c7d..49928c0d34dc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -5987,14 +5987,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=}" 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 1d36ebcd1221..7fb6ab4293a0 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -35,6 +35,7 @@ from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.spec_info import SpeculativeAlgorithm 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 CustomTestCase register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -65,29 +66,17 @@ def _serving(self, version: str) -> _ServingStub: 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] + running, queued, chunked, staged = (object() for _ in range(4)) + scheduler = Scheduler.__new__(Scheduler) + empty_request_lists(scheduler) + scheduler.running_reqs.join([running]) + scheduler.waiting_queue = [queued] + scheduler.chunked_req = chunked + scheduler.hisparse_coordinator = SimpleNamespace( + request_lists=lambda: {ListKind.HISPARSE_STAGING: [staged]} ) with patch( @@ -97,7 +86,7 @@ def test_every_source_of_live_requests_is_stamped(self): Scheduler.record_weight_version_change(scheduler, new_version="v2") self.assertEqual( - set(recorder.call_args.args[0]), {inflight, queued, chunked, staged} + set(recorder.call_args.args[0]), {running, queued, chunked, staged} ) diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 37adfe50e95f..7b2bb4ee2833 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -1,15 +1,11 @@ import random import unittest -from collections import deque from types import SimpleNamespace from unittest.mock import patch -from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.deferred_responses import ( - DeferredResponses, -) +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.managers.utils import make_abort_req from sglang.srt.utils.weight_versions import ( WeightVersionEvent, @@ -22,11 +18,8 @@ weight_version_spans_to_json, ) from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import ( - CustomTestCase, - enter_scope, - published_topology, -) +from sglang.test.scheduler_test_utils import empty_request_lists +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -302,38 +295,14 @@ def override(self, source, **fields): self.serving.weight_version = fields["weight_version"] -class _SchedulerStub: - _collect_inflight_batches = Scheduler._collect_inflight_batches - collect_inflight_reqs = Scheduler.collect_inflight_reqs - - def __init__( - self, - version, - running, - waiting, - chunked=None, - last_batch=None, - hisparse=None, - ): - self.serving = _ServingStub(version) - self.enable_continuous_input_polling = False - self.result_queue = deque() - self.running_batch = SimpleNamespace(reqs=running) - self.last_batch = last_batch - self.waiting_queue = waiting - self.chunked_req = chunked - self.hisparse_coordinator = hisparse - self.deferred_responses = DeferredResponses( - disaggregation_mode=DisaggregationMode.NULL, - output_streamer=None, - send_abort=None, - ) - - class TestSchedulerRecordWeightVersionChange(CustomTestCase): - def _scheduler(self, *args, pp_size=1, **kwargs): - enter_scope(self, published_topology(pp_size=pp_size)) - scheduler = _SchedulerStub(*args, **kwargs) + def _scheduler(self, version, running=(), waiting=(), chunked=None): + scheduler = Scheduler.__new__(Scheduler) + empty_request_lists(scheduler) + scheduler.running_reqs.join(running) + scheduler.waiting_queue = list(waiting) + scheduler.chunked_req = chunked + scheduler.serving = _ServingStub(version) for name, value in ( ("get_serving", scheduler.serving), ("get_context", _ContextStub(scheduler.serving)), @@ -364,44 +333,27 @@ def test_records_on_running_waiting_and_chunked_requests(self): self.assertEqual(req.weight_version_events[0].old_version, "v1") self.assertEqual(req.weight_version_events[0].num_output_tokens, output_len) - def test_records_on_last_batch_requests_exactly_once(self): - """A request in both the running and last batch is recorded once, and last-batch-only requests are recorded too.""" - shared_req = _ReqStub(3) - prefill_req = _ReqStub(1) - scheduler = self._scheduler( - "v1", - [shared_req], - [], - last_batch=SimpleNamespace(reqs=[shared_req, prefill_req]), - ) + def test_records_on_requests_outside_running_batches(self): + """Requests waiting in the grammar or PD queues hold output across the update too.""" + grammar_req = _ReqStub(2) + sending_req = _ReqStub(4) + scheduler = self._scheduler("v1") + scheduler.grammar_manager.grammar_queue.join([grammar_req]) + scheduler.disagg_prefill_transfer_queue.join([sending_req]) Scheduler.record_weight_version_change(scheduler, new_version="v2") - self.assertEqual(len(shared_req.weight_version_events), 1) - self.assertEqual(len(prefill_req.weight_version_events), 1) - - def test_records_on_all_pp_microbatches(self): - """With pipeline parallelism every microbatch is visited, not just the selected one.""" - mb0_req = _ReqStub(2) - mb1_req = _ReqStub(4) - pending_req = _ReqStub(6) - scheduler = self._scheduler("v1", [], [], pp_size=2) - scheduler.running_mbs = [ - SimpleNamespace(reqs=[mb0_req]), - SimpleNamespace(reqs=[mb1_req]), - ] - scheduler.mbs = [None, SimpleNamespace(reqs=[pending_req])] - - Scheduler.record_weight_version_change(scheduler, new_version="v2") - - for req in (mb0_req, mb1_req, pending_req): + for req in (grammar_req, sending_req): self.assertEqual(len(req.weight_version_events), 1) + self.assertEqual(req.weight_version_events[0].old_version, "v1") def test_records_on_hisparse_staging_requests(self): """Requests parked in the HiSparse staging queue are swept like any live request.""" staging_req = _ReqStub(2) - coordinator = SimpleNamespace(ack_staging_queue=[staging_req]) - scheduler = self._scheduler("v1", [], [], hisparse=coordinator) + scheduler = self._scheduler("v1") + scheduler.hisparse_coordinator = SimpleNamespace( + request_lists=lambda: {ListKind.HISPARSE_STAGING: [staging_req]} + ) Scheduler.record_weight_version_change(scheduler, new_version="v2") @@ -412,7 +364,7 @@ def test_records_on_hisparse_staging_requests(self): def test_same_version_is_a_noop(self): """Re-announcing the current version must not record events.""" req = _ReqStub(3) - scheduler = self._scheduler("v1", [req], []) + scheduler = self._scheduler("v1", [req]) Scheduler.record_weight_version_change(scheduler, new_version="v1") @@ -422,7 +374,7 @@ def test_same_version_is_a_noop(self): def test_none_version_is_a_noop(self): """An update without a weight version must not disturb attribution.""" req = _ReqStub(3) - scheduler = self._scheduler("v1", [req], []) + scheduler = self._scheduler("v1", [req]) Scheduler.record_weight_version_change(scheduler, new_version=None) From 0cc8727965afde658f5e205f250064cc5abded6d Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 22:29:23 -0800 Subject: [PATCH 2/5] build the held-token weight version fixture from empty request lists --- .../unit/managers/test_deferred_output.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/test/registered/unit/managers/test_deferred_output.py b/test/registered/unit/managers/test_deferred_output.py index 6527d0f90f0d..85b99226c3a9 100644 --- a/test/registered/unit/managers/test_deferred_output.py +++ b/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) From 6fd56746319950fe1df666dce4a203f3325f382a Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 22:29:52 -0800 Subject: [PATCH 3/5] read queue gauges and load from request lists by kind --- python/sglang/srt/managers/scheduler.py | 8 +-- .../scheduler_components/load_inquirer.py | 28 +++++----- .../scheduler_components/metrics_reporter.py | 54 +++++++------------ .../unit/managers/test_load_inquirer.py | 5 +- .../test_forward_pass_metrics.py | 5 +- .../observability/test_stat_loggers_di.py | 9 ++-- 6 files changed, 38 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 49928c0d34dc..8ad5ca3b1496 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 ), diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 585ca9008470..00998b021c46 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 3a442a15dd4a..3b15f6c96c31 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/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 diff --git a/test/registered/unit/managers/test_load_inquirer.py b/test/registered/unit/managers/test_load_inquirer.py index a77980b99153..719f5b6bdf15 100644 --- a/test/registered/unit/managers/test_load_inquirer.py +++ b/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, @@ -238,6 +236,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) diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index 1ed15fc950ab..2c81124aee98 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/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) ) diff --git a/test/registered/unit/observability/test_stat_loggers_di.py b/test/registered/unit/observability/test_stat_loggers_di.py index 6f795472ba01..4d306e50213e 100644 --- a/test/registered/unit/observability/test_stat_loggers_di.py +++ b/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( From e5f9f6b9530117f35d5631b5d5af97c1d77c2e02 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 2 Nov 2026 23:14:05 -0800 Subject: [PATCH 4/5] feed the prealloc ready-count test through request lists --- test/registered/unit/managers/test_load_inquirer.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/test/registered/unit/managers/test_load_inquirer.py b/test/registered/unit/managers/test_load_inquirer.py index 719f5b6bdf15..65f4551ae135 100644 --- a/test/registered/unit/managers/test_load_inquirer.py +++ b/test/registered/unit/managers/test_load_inquirer.py @@ -134,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,), From 02ff2fcf71a5e56f80e7a8ecafe5bd542aa470c6 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 3 Nov 2026 00:15:24 -0800 Subject: [PATCH 5/5] merge the weight version list coverage tests into one --- .../test_scheduler_weight_version_tracking.py | 45 ------------------- .../unit/utils/test_weight_versions.py | 40 ++++------------- 2 files changed, 9 insertions(+), 76 deletions(-) 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 7fb6ab4293a0..f453ef6d7d06 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -35,7 +35,6 @@ from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.spec_info import SpeculativeAlgorithm 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 CustomTestCase register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -46,50 +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 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") - running, queued, chunked, staged = (object() for _ in range(4)) - scheduler = Scheduler.__new__(Scheduler) - empty_request_lists(scheduler) - scheduler.running_reqs.join([running]) - scheduler.waiting_queue = [queued] - scheduler.chunked_req = chunked - scheduler.hisparse_coordinator = SimpleNamespace( - request_lists=lambda: {ListKind.HISPARSE_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]), {running, 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.""" diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 7b2bb4ee2833..b3cc3ef4e398 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -5,7 +5,6 @@ from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.managers.utils import make_abort_req from sglang.srt.utils.weight_versions import ( WeightVersionEvent, @@ -312,14 +311,19 @@ def _scheduler(self, version, running=(), waiting=(), chunked=None): self.addCleanup(patcher.stop) return scheduler - def test_records_on_running_waiting_and_chunked_requests(self): - """A version change records an event on every live request holding output tokens.""" + def test_records_on_every_request_list(self): + """A version change records an event on every live request holding output + tokens, whichever list holds it.""" running_req = _ReqStub(3) waiting_req = _ReqStub(5) chunked_req = _ReqStub(2) + grammar_req = _ReqStub(4) + sending_req = _ReqStub(6) scheduler = self._scheduler( "v1", [running_req], [waiting_req], chunked=chunked_req ) + scheduler.grammar_manager.grammar_queue.join([grammar_req]) + scheduler.disagg_prefill_transfer_queue.join([sending_req]) Scheduler.record_weight_version_change(scheduler, new_version="v2") @@ -328,39 +332,13 @@ def test_records_on_running_waiting_and_chunked_requests(self): (running_req, 3), (waiting_req, 5), (chunked_req, 2), + (grammar_req, 4), + (sending_req, 6), ): self.assertEqual(len(req.weight_version_events), 1) self.assertEqual(req.weight_version_events[0].old_version, "v1") self.assertEqual(req.weight_version_events[0].num_output_tokens, output_len) - def test_records_on_requests_outside_running_batches(self): - """Requests waiting in the grammar or PD queues hold output across the update too.""" - grammar_req = _ReqStub(2) - sending_req = _ReqStub(4) - scheduler = self._scheduler("v1") - scheduler.grammar_manager.grammar_queue.join([grammar_req]) - scheduler.disagg_prefill_transfer_queue.join([sending_req]) - - Scheduler.record_weight_version_change(scheduler, new_version="v2") - - for req in (grammar_req, sending_req): - self.assertEqual(len(req.weight_version_events), 1) - self.assertEqual(req.weight_version_events[0].old_version, "v1") - - def test_records_on_hisparse_staging_requests(self): - """Requests parked in the HiSparse staging queue are swept like any live request.""" - staging_req = _ReqStub(2) - scheduler = self._scheduler("v1") - scheduler.hisparse_coordinator = SimpleNamespace( - request_lists=lambda: {ListKind.HISPARSE_STAGING: [staging_req]} - ) - - Scheduler.record_weight_version_change(scheduler, new_version="v2") - - self.assertEqual(len(staging_req.weight_version_events), 1) - self.assertEqual(staging_req.weight_version_events[0].old_version, "v1") - self.assertEqual(staging_req.weight_version_events[0].num_output_tokens, 2) - def test_same_version_is_a_noop(self): """Re-announcing the current version must not record events.""" req = _ReqStub(3)